From 2567efa39fdc52b37deadad383455b5adcd7e9bb Mon Sep 17 00:00:00 2001 From: Anthony Alaribe Date: Thu, 29 Jan 2026 23:12:10 -0800 Subject: [PATCH 1/7] Add Variant type support and fix SLT tests - Add Parquet Variant binary encoding support via datafusion-variant crate - Implement VariantAwareExprPlanner for -> and ->> operators on Variant columns - Add jsonb_path_exists UDF for JSONPath queries on Variant/JSON columns - Register variant functions: json_to_variant, variant_to_json, variant_get, etc. - Add is_variant_type helper in schema_loader for Variant type detection - Update schema with Variant columns: context, events, links, attributes, resource - Fix time_bucket UDF to handle Utf8/Utf8View/LargeUtf8 string types - Fix SLT tests: correct status_message type (string not array) - Fix json_functions.slt to use json_to_variant for Variant column inserts - Add variant_functions.slt tests for round-trip, path extraction, arrow operators --- Cargo.lock | 197 +++++++++-- Cargo.toml | 10 +- schemas/otel_logs_and_spans.yaml | 10 +- src/database.rs | 90 ++++- src/functions.rs | 405 ++++++++++++++++++++++- src/schema_loader.rs | 26 ++ tests/connection_pressure_test.rs | 5 +- tests/integration_test.rs | 5 +- tests/slt/custom_functions.slt | 8 +- tests/slt/function_availability_test.slt | 4 +- tests/slt/integration.slt | 2 +- tests/slt/json_functions.slt | 4 +- tests/slt/variant_functions.slt | 269 +++++++++++++++ tests/sqllogictest.rs | 5 +- 14 files changed, 979 insertions(+), 61 deletions(-) create mode 100644 tests/slt/variant_functions.slt diff --git a/Cargo.lock b/Cargo.lock index 38080d0c..1c01d908 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -172,9 +172,9 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb372a7cbcac02a35d3fb7b3fc1f969ec078e871f9bb899bf00a2e1809bec8a3" +checksum = "2a2b10dcb159faf30d3f81f6d56c1211a5bea2ca424eabe477648a44b993320e" dependencies = [ "arrow-arith", "arrow-array", @@ -193,9 +193,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f377dcd19e440174596d83deb49cd724886d91060c07fec4f67014ef9d54049" +checksum = "288015089e7931843c80ed4032c5274f02b37bcb720c4a42096d50b390e70372" dependencies = [ "arrow-array", "arrow-buffer", @@ -207,9 +207,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a23eaff85a44e9fa914660fb0d0bb00b79c4a3d888b5334adb3ea4330c84f002" +checksum = "65ca404ea6191e06bf30956394173337fa9c35f445bd447fe6c21ab944e1a23c" dependencies = [ "ahash 0.8.12", "arrow-buffer", @@ -226,9 +226,9 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2819d893750cb3380ab31ebdc8c68874dd4429f90fd09180f3c93538bd21626" +checksum = "36356383099be0151dacc4245309895f16ba7917d79bdb71a7148659c9206c56" dependencies = [ "bytes", "half", @@ -238,9 +238,9 @@ dependencies = [ [[package]] name = "arrow-cast" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3d131abb183f80c450d4591dc784f8d7750c50c6e2bc3fcaad148afc8361271" +checksum = "9c8e372ed52bd4ee88cc1e6c3859aa7ecea204158ac640b10e187936e7e87074" dependencies = [ "arrow-array", "arrow-buffer", @@ -260,9 +260,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2275877a0e5e7e7c76954669366c2aa1a829e340ab1f612e647507860906fb6b" +checksum = "8e4100b729fe656f2e4fb32bc5884f14acf9118d4ad532b7b33c1132e4dce896" dependencies = [ "arrow-array", "arrow-cast", @@ -275,9 +275,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05738f3d42cb922b9096f7786f606fcb8669260c2640df8490533bb2fa38c9d3" +checksum = "bf87f4ff5fc13290aa47e499a8b669a82c5977c6a1fedce22c7f542c1fd5a597" dependencies = [ "arrow-buffer", "arrow-schema", @@ -288,9 +288,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d09446e8076c4b3f235603d9ea7c5494e73d441b01cd61fb33d7254c11964b3" +checksum = "eb3ca63edd2073fcb42ba112f8ae165df1de935627ead6e203d07c99445f2081" dependencies = [ "arrow-array", "arrow-buffer", @@ -304,9 +304,9 @@ dependencies = [ [[package]] name = "arrow-json" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "371ffd66fa77f71d7628c63f209c9ca5341081051aa32f9c8020feb0def787c0" +checksum = "a36b2332559d3310ebe3e173f75b29989b4412df4029a26a30cc3f7da0869297" dependencies = [ "arrow-array", "arrow-buffer", @@ -328,9 +328,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc94fc7adec5d1ba9e8cd1b1e8d6f72423b33fe978bf1f46d970fafab787521" +checksum = "13c4e0530272ca755d6814218dffd04425c5b7854b87fa741d5ff848bf50aa39" dependencies = [ "arrow-array", "arrow-buffer", @@ -356,9 +356,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "169676f317157dc079cc5def6354d16db63d8861d61046d2f3883268ced6f99f" +checksum = "b07f52788744cc71c4628567ad834cadbaeb9f09026ff1d7a4120f69edf7abd3" dependencies = [ "arrow-array", "arrow-buffer", @@ -369,9 +369,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d27609cd7dd45f006abae27995c2729ef6f4b9361cde1ddd019dc31a5aa017e0" +checksum = "6bb63203e8e0e54b288d0d8043ca8fa1013820822a27692ef1b78a977d879f2c" dependencies = [ "bitflags", "serde", @@ -381,9 +381,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae980d021879ea119dd6e2a13912d81e64abed372d53163e804dfe84639d8010" +checksum = "c96d8a1c180b44ecf2e66c9a2f2bbcb8b1b6f14e165ce46ac8bde211a363411b" dependencies = [ "ahash 0.8.12", "arrow-array", @@ -395,9 +395,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "57.1.0" +version = "57.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf35e8ef49dcf0c5f6d175edee6b8af7b45611805333129c541a8b89a0fc0534" +checksum = "a8ad6a81add9d3ea30bf8374ee8329992c7fd246ffd8b7e2f48a3cea5aa0cc9a" dependencies = [ "arrow-array", "arrow-buffer", @@ -2518,6 +2518,18 @@ dependencies = [ "unicode-width 0.2.2", ] +[[package]] +name = "datafusion-variant" +version = "0.1.0" +dependencies = [ + "arrow", + "arrow-schema", + "datafusion", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", +] + [[package]] name = "delegate" version = "0.13.5" @@ -2574,7 +2586,6 @@ dependencies = [ [[package]] name = "deltalake" version = "0.30.1" -source = "git+https://github.com/delta-io/delta-rs.git?rev=ffb794ba0745394fc4b747a4ef2e11c2d4ec086a#ffb794ba0745394fc4b747a4ef2e11c2d4ec086a" dependencies = [ "ctor", "delta_kernel", @@ -2585,7 +2596,6 @@ dependencies = [ [[package]] name = "deltalake-aws" version = "0.13.0" -source = "git+https://github.com/delta-io/delta-rs.git?rev=ffb794ba0745394fc4b747a4ef2e11c2d4ec086a#ffb794ba0745394fc4b747a4ef2e11c2d4ec086a" dependencies = [ "async-trait", "aws-config", @@ -2611,7 +2621,6 @@ dependencies = [ [[package]] name = "deltalake-core" version = "0.30.1" -source = "git+https://github.com/delta-io/delta-rs.git?rev=ffb794ba0745394fc4b747a4ef2e11c2d4ec086a#ffb794ba0745394fc4b747a4ef2e11c2d4ec086a" dependencies = [ "arrow", "arrow-arith", @@ -2664,7 +2673,6 @@ dependencies = [ [[package]] name = "deltalake-derive" version = "0.30.0" -source = "git+https://github.com/delta-io/delta-rs.git?rev=ffb794ba0745394fc4b747a4ef2e11c2d4ec086a#ffb794ba0745394fc4b747a4ef2e11c2d4ec086a" dependencies = [ "convert_case", "itertools 0.14.0", @@ -3944,6 +3952,15 @@ version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" +[[package]] +name = "inventory" +version = "0.3.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc61209c082fbeb19919bee74b176221b27223e27b65d781eb91af24eb1fb46e" +dependencies = [ + "rustversion", +] + [[package]] name = "io-uring" version = "0.7.11" @@ -4409,6 +4426,12 @@ dependencies = [ "autocfg", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -4455,6 +4478,16 @@ dependencies = [ "getrandom 0.2.16", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -4816,6 +4849,50 @@ dependencies = [ "zstd", ] +[[package]] +name = "parquet-variant" +version = "57.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c254fac16af78ad96aa442290cb6504951c4d484fdfcfe58f4588033d30e4c8f" +dependencies = [ + "arrow-schema", + "chrono", + "half", + "indexmap 2.12.1", + "simdutf8", + "uuid", +] + +[[package]] +name = "parquet-variant-compute" +version = "57.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2178772f1c5ad7e5da8b569d986d3f5cbb4a4cee915925f28fdc700dbb2e80cf" +dependencies = [ + "arrow", + "arrow-schema", + "chrono", + "half", + "indexmap 2.12.1", + "parquet-variant", + "parquet-variant-json", + "uuid", +] + +[[package]] +name = "parquet-variant-json" +version = "57.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a1510daa121c04848368f9c38d0be425b9418c70be610ecc0aa8071738c0ef3" +dependencies = [ + "arrow-schema", + "base64", + "chrono", + "parquet-variant", + "serde_json", + "uuid", +] + [[package]] name = "paste" version = "1.0.15" @@ -5968,6 +6045,56 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_json_path" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b992cea3194eea663ba99a042d61cea4bd1872da37021af56f6a37e0359b9d33" +dependencies = [ + "inventory", + "nom", + "regex", + "serde", + "serde_json", + "serde_json_path_core", + "serde_json_path_macros", + "thiserror", +] + +[[package]] +name = "serde_json_path_core" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dde67d8dfe7d4967b5a95e247d4148368ddd1e753e500adb34b3ffe40c6bc1bc" +dependencies = [ + "inventory", + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "serde_json_path_macros" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "517acfa7f77ddaf5c43d5f119c44a683774e130b4247b7d3210f8924506cfac8" +dependencies = [ + "inventory", + "serde_json_path_core", + "serde_json_path_macros_internal", +] + +[[package]] +name = "serde_json_path_macros_internal" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aafbefbe175fa9bf03ca83ef89beecff7d2a95aaacd5732325b90ac8c3bd7b90" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.114", +] + [[package]] name = "serde_spanned" version = "1.0.4" @@ -6746,6 +6873,7 @@ dependencies = [ "aws-sdk-dynamodb", "aws-sdk-s3", "aws-types", + "base64", "bincode 2.0.1", "bytes", "chrono", @@ -6758,6 +6886,7 @@ dependencies = [ "datafusion-functions-json", "datafusion-postgres", "datafusion-tracing", + "datafusion-variant", "delta_kernel", "deltalake", "dotenv", @@ -6773,6 +6902,9 @@ dependencies = [ "opentelemetry-otlp", "opentelemetry_sdk", "parking_lot", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", "rand 0.9.2", "regex", "scopeguard", @@ -6780,6 +6912,7 @@ dependencies = [ "serde_arrow", "serde_bytes", "serde_json", + "serde_json_path", "serde_with", "serde_yaml", "serial_test", diff --git a/Cargo.toml b/Cargo.toml index fe2ba79e..a7596e4b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 = [ +# Updated to delta-rs with Variant type support (adds variantType feature detection) +deltalake = { path = "/Users/tonyalaribe/Projects/apitoolkit/datafusion-projects/delta-rs/crates/deltalake", 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 = { path = "/Users/tonyalaribe/Projects/apitoolkit/datafusion-projects/datafusion-variant" } +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" } diff --git a/schemas/otel_logs_and_spans.yaml b/schemas/otel_logs_and_spans.yaml index bcb0aa46..3cf42bee 100644 --- a/schemas/otel_logs_and_spans.yaml +++ b/schemas/otel_logs_and_spans.yaml @@ -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 diff --git a/src/database.rs b/src/database.rs index c025dad2..b61f620c 100644 --- a/src/database.rs +++ b/src/database.rs @@ -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; @@ -82,6 +82,77 @@ pub fn extract_project_id(batch: &RecordBatch) -> Option { }) } +/// 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 { + use datafusion::arrow::array::{ArrayRef, LargeStringArray, StringArray, StringViewArray}; + use datafusion::arrow::datatypes::{DataType, Field}; + + let batch_schema = batch.schema(); + let mut columns: Vec = batch.columns().to_vec(); + let mut new_fields: Vec> = 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; + } + 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 = match col_type { + DataType::Utf8View => { + let arr = col.as_any().downcast_ref::().unwrap(); + Some(Arc::new(json_strings_to_variant(arr.iter()))) + } + DataType::Utf8 => { + let arr = col.as_any().downcast_ref::().unwrap(); + Some(Arc::new(json_strings_to_variant(arr.iter()))) + } + DataType::LargeUtf8 => { + let arr = col.as_any().downcast_ref::().unwrap(); + 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 +fn json_strings_to_variant<'a>(iter: impl Iterator>) -> 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 item in items { + match item { + Some(json_str) => { + if let Err(e) = builder.append_json(json_str) { + warn!("Failed to parse JSON '{}': {}, inserting as null", json_str, e); + builder.append_null(); + } + } + None => builder.append_null(), + } + } + + builder.build().into() +} + // Compression level for parquet files - kept for WriterProperties fallback const ZSTD_COMPRESSION_LEVEL: i32 = 3; @@ -712,11 +783,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 + // Register custom PostgreSQL-compatible functions BEFORE JSON functions + // so VariantAwareExprPlanner gets first chance at -> and ->> operators crate::functions::register_custom_functions(ctx).map_err(|e| DataFusionError::Execution(format!("Failed to register custom functions: {}", e)))?; + // JSON functions (includes JsonExprPlanner for -> and ->> on string columns) + self.register_json_functions(ctx); + Ok(()) } @@ -1892,14 +1966,18 @@ impl DataSink for ProjectRoutingTable { let span = tracing::Span::current(); let mut total_row_count = 0; let mut project_batches: HashMap> = 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); diff --git a/src/functions.rs b/src/functions.rs index c12a2390..60eef686 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -6,18 +6,188 @@ use datafusion::arrow::array::{ TimestampMicrosecondArray, TimestampNanosecondArray, }; use datafusion::arrow::datatypes::{DataType, TimeUnit}; -use datafusion::common::{DataFusionError, ScalarValue, not_impl_err}; +use datafusion::common::{DFSchema, DataFusionError, ExprSchema, ScalarValue, not_impl_err}; +use datafusion::logical_expr::ExprSchemable; use datafusion::logical_expr::{ - Accumulator, AggregateUDF, ColumnarValue, ScalarFunctionArgs, ScalarFunctionImplementation, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, Volatility, - create_udaf, create_udf, + Accumulator, AggregateUDF, ColumnarValue, Expr, ScalarFunctionArgs, ScalarFunctionImplementation, + ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, Volatility, create_udaf, create_udf, + expr::{Alias, ScalarFunction}, + planner::{ExprPlanner, PlannerResult, RawBinaryExpr}, }; +use datafusion::sql::sqlparser::ast::BinaryOperator; use serde_json::{Value as JsonValue, json}; use std::any::Any; use std::sync::Arc; use tdigests::TDigest; +use crate::schema_loader::is_variant_type; + +// ============================================================================ +// Variant-Aware Expression Planner +// ============================================================================ + +/// ExprPlanner that intercepts -> and ->> operators on Variant columns +/// and rewrites them to efficient variant_get calls with flattened dot-paths. +#[derive(Debug, Default)] +pub struct VariantAwareExprPlanner; + +/// Path component for building variant_get paths +#[derive(Debug, Clone)] +enum PathComponent { + Field(String), + Index(i64), +} + +impl ExprPlanner for VariantAwareExprPlanner { + fn plan_binary_op(&self, expr: RawBinaryExpr, schema: &DFSchema) -> datafusion::error::Result> { + let is_long_arrow = match &expr.op { + BinaryOperator::Arrow => false, + BinaryOperator::LongArrow => true, + _ => return Ok(PlannerResult::Original(expr)), + }; + + // Recursively collect path components from chained operators + let (base_expr, mut path_parts) = collect_arrow_chain(&expr.left); + if let Some(component) = extract_path_component(&expr.right) { + path_parts.push(component); + } else { + return Ok(PlannerResult::Original(expr)); + } + + // Check if base column is Variant type + if !is_variant_column(&base_expr, schema) { + return Ok(PlannerResult::Original(expr)); // Let JSON planner handle + } + + // Build dot-path: ["user", "name"] → "user.name", ["items", Index(0)] → "items[0]" + let full_path = build_variant_path(&path_parts); + + // Create variant_get function call + let variant_get_udf = ScalarUDF::from(datafusion_variant::VariantGetUdf::default()); + let path_literal = Expr::Literal(ScalarValue::Utf8(Some(full_path.clone())), None); + let variant_get_call = Expr::ScalarFunction(ScalarFunction { + func: Arc::new(variant_get_udf), + args: vec![base_expr.clone(), path_literal], + }); + + // For ->> wrap with variant_to_json for text output + let result = if is_long_arrow { + let variant_to_json_udf = ScalarUDF::from(datafusion_variant::VariantToJsonUdf::default()); + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(variant_to_json_udf), + args: vec![variant_get_call], + }) + } else { + variant_get_call + }; + + // Create alias to preserve original SQL representation + let op_str = if is_long_arrow { "->>" } else { "->" }; + let alias_name = format!("{} {} {}", expr_repr(&base_expr), op_str, path_repr(&path_parts)); + Ok(PlannerResult::Planned(Expr::Alias(Alias::new(result, None::<&str>, alias_name)))) + } +} + +/// Recursively collect chained arrow expressions into base + path components +fn collect_arrow_chain(expr: &Expr) -> (Expr, Vec) { + match expr { + Expr::BinaryExpr(binary) if matches!(binary.op, datafusion::logical_expr::Operator::Arrow) => { + let (base, mut parts) = collect_arrow_chain(&binary.left); + if let Some(component) = extract_path_component(&binary.right) { + parts.push(component); + } + (base, parts) + } + Expr::Alias(alias) => collect_arrow_chain(&alias.expr), + _ => (expr.clone(), vec![]), + } +} + +/// Extract path component from expression (string literal or integer) +fn extract_path_component(expr: &Expr) -> Option { + match expr { + Expr::Literal(ScalarValue::Utf8(Some(s)), _) => Some(PathComponent::Field(s.clone())), + Expr::Literal(ScalarValue::Utf8View(Some(s)), _) => Some(PathComponent::Field(s.clone())), + Expr::Literal(ScalarValue::LargeUtf8(Some(s)), _) => Some(PathComponent::Field(s.clone())), + Expr::Literal(ScalarValue::Int64(Some(i)), _) => Some(PathComponent::Index(*i)), + Expr::Literal(ScalarValue::Int32(Some(i)), _) => Some(PathComponent::Index(*i as i64)), + Expr::Literal(ScalarValue::UInt64(Some(i)), _) => Some(PathComponent::Index(*i as i64)), + Expr::Literal(ScalarValue::UInt32(Some(i)), _) => Some(PathComponent::Index(*i as i64)), + _ => None, + } +} + +/// Check if expression evaluates to a Variant type +fn is_variant_column(expr: &Expr, schema: &DFSchema) -> bool { + match expr { + // Direct column reference + Expr::Column(col) => schema + .field_from_column(col) + .map(|f| is_variant_type(f.data_type())) + .unwrap_or(false), + // Unwrap aliases + Expr::Alias(alias) => is_variant_column(&alias.expr, schema), + // Check if it's a call to a variant-producing function + Expr::ScalarFunction(func) => { + let name = func.func.name(); + matches!(name, "json_to_variant" | "variant_get" | "cast_to_variant" + | "variant_object_construct" | "variant_list_construct" + | "variant_object_insert" | "variant_list_insert") + } + // Try to get the type for other expressions + _ => expr.get_type(schema) + .map(|dt| is_variant_type(&dt)) + .unwrap_or(false), + } +} + +/// Build variant_get path string from components +fn build_variant_path(parts: &[PathComponent]) -> String { + let mut path = String::new(); + for (i, part) in parts.iter().enumerate() { + match part { + PathComponent::Field(name) => { + if i > 0 { + path.push('.'); + } + path.push_str(name); + } + PathComponent::Index(idx) => { + path.push('['); + path.push_str(&idx.to_string()); + path.push(']'); + } + } + } + path +} + +/// Generate SQL-like representation for expression (for alias) +fn expr_repr(expr: &Expr) -> String { + match expr { + Expr::Column(col) => col.name.clone(), + Expr::Alias(alias) => alias.name.clone(), + _ => "expr".to_string(), + } +} + +/// Generate path representation for alias +fn path_repr(parts: &[PathComponent]) -> String { + parts + .iter() + .map(|p| match p { + PathComponent::Field(s) => format!("'{}'", s), + PathComponent::Index(i) => i.to_string(), + }) + .collect::>() + .join("->") +} + /// Register all custom PostgreSQL-compatible functions pub fn register_custom_functions(ctx: &mut datafusion::execution::context::SessionContext) -> Result<()> { + // Register Variant-aware expr planner (must be before JSON planner for priority) + datafusion::execution::FunctionRegistry::register_expr_planner(ctx, Arc::new(VariantAwareExprPlanner))?; + // Register to_char function ctx.register_udf(create_to_char_udf()); @@ -45,6 +215,21 @@ pub fn register_custom_functions(ctx: &mut datafusion::execution::context::Sessi // Register approx_percentile scalar function ctx.register_udf(create_approx_percentile_udf()); + // Register variant functions from datafusion-variant + ctx.register_udf(ScalarUDF::from(datafusion_variant::JsonToVariantUdf::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::VariantToJsonUdf::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::VariantGetUdf::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::CastToVariantUdf::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::IsVariantNullUdf::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::VariantPretty::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::VariantListConstruct::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::VariantListInsert::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::VariantObjectConstruct::default())); + ctx.register_udf(ScalarUDF::from(datafusion_variant::VariantObjectInsert::default())); + + // Register jsonb_path_exists for JSONPath queries on Variant columns + ctx.register_udf(create_jsonb_path_exists_udf()); + Ok(()) } @@ -684,7 +869,9 @@ fn create_time_bucket_udf() -> ScalarUDF { // Extract interval string let interval_str = match &args[0] { ColumnarValue::Scalar(scalar) => match scalar { - datafusion::scalar::ScalarValue::Utf8(Some(s)) => s.clone(), + datafusion::scalar::ScalarValue::Utf8(Some(s)) + | datafusion::scalar::ScalarValue::Utf8View(Some(s)) + | datafusion::scalar::ScalarValue::LargeUtf8(Some(s)) => s.clone(), _ => return Err(DataFusionError::Execution("Interval must be a UTF8 string".to_string())), }, ColumnarValue::Array(_) => { @@ -1022,6 +1209,216 @@ impl ScalarUDFImpl for ApproxPercentileUDF { } } +// ============================================================================ +// jsonb_path_exists UDF for JSONPath queries on Variant/JSON columns +// ============================================================================ + +/// Create the jsonb_path_exists UDF for PostgreSQL-compatible JSONPath queries +fn create_jsonb_path_exists_udf() -> ScalarUDF { + ScalarUDF::from(JsonbPathExistsUDF::new()) +} + +#[derive(Debug, Hash, Eq, PartialEq)] +struct JsonbPathExistsUDF { + signature: Signature, +} + +impl JsonbPathExistsUDF { + fn new() -> Self { + Self { + // Accept Variant struct or JSON string as first arg, path string as second + signature: Signature::any(2, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for JsonbPathExistsUDF { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "jsonb_path_exists" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> datafusion::error::Result { + if args.args.len() != 2 { + return Err(DataFusionError::Execution( + "jsonb_path_exists requires exactly 2 arguments: json/variant and jsonpath".to_string(), + )); + } + + let json_array = match &args.args[0] { + ColumnarValue::Array(array) => array.clone(), + ColumnarValue::Scalar(scalar) => scalar.to_array()?, + }; + + let path_str = match &args.args[1] { + ColumnarValue::Scalar(scalar) => match scalar { + ScalarValue::Utf8(Some(s)) | ScalarValue::Utf8View(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => s.clone(), + _ => return Err(DataFusionError::Execution("JSONPath must be a string".to_string())), + }, + ColumnarValue::Array(_) => { + return Err(DataFusionError::Execution("JSONPath must be a scalar string".to_string())); + } + }; + + // Parse the JSONPath expression + let json_path = serde_json_path::JsonPath::parse(&path_str) + .map_err(|e| DataFusionError::Execution(format!("Invalid JSONPath: {}", e)))?; + + // Process based on input type + let result = if is_variant_type(json_array.data_type()) { + // Handle Variant struct type + evaluate_jsonpath_on_variant(&json_array, &json_path)? + } else { + // Handle JSON string type + evaluate_jsonpath_on_json_string(&json_array, &json_path)? + }; + + Ok(ColumnarValue::Array(result)) + } +} + +/// Convert parquet_variant::Variant to serde_json::Value +fn variant_to_serde_json(variant: &parquet_variant::Variant) -> JsonValue { + use base64::Engine; + use parquet_variant::Variant; + + match variant { + Variant::Null => JsonValue::Null, + Variant::BooleanTrue => JsonValue::Bool(true), + Variant::BooleanFalse => JsonValue::Bool(false), + Variant::Int8(v) => json!(*v), + Variant::Int16(v) => json!(*v), + Variant::Int32(v) => json!(*v), + Variant::Int64(v) => json!(*v), + Variant::Float(v) => json!(*v), + Variant::Double(v) => json!(*v), + Variant::Decimal4(d) => json!(d.to_string()), + Variant::Decimal8(d) => json!(d.to_string()), + Variant::Decimal16(d) => json!(d.to_string()), + Variant::Date(v) => json!(*v), + Variant::Time(v) => json!(*v), + Variant::Uuid(v) => json!(v.to_string()), + Variant::TimestampMicros(v) => json!(*v), + Variant::TimestampNtzMicros(v) => json!(*v), + Variant::TimestampNanos(v) => json!(*v), + Variant::TimestampNtzNanos(v) => json!(*v), + Variant::Binary(bytes) => json!(base64::engine::general_purpose::STANDARD.encode(bytes)), + Variant::String(s) => JsonValue::String(s.to_string()), + Variant::ShortString(s) => JsonValue::String(s.as_str().to_string()), + Variant::Object(obj) => { + let mut map = serde_json::Map::new(); + for (key, value) in obj.iter() { + map.insert(key.to_string(), variant_to_serde_json(&value)); + } + JsonValue::Object(map) + } + Variant::List(list) => { + let items: Vec = list.iter().map(|v| variant_to_serde_json(&v)).collect(); + JsonValue::Array(items) + } + } +} + +/// Evaluate JSONPath on a Variant (Struct) array +fn evaluate_jsonpath_on_variant(array: &ArrayRef, json_path: &serde_json_path::JsonPath) -> datafusion::error::Result { + use datafusion::arrow::array::StructArray; + use parquet_variant::Variant; + + let struct_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Execution("Expected Variant struct array".to_string()))?; + + let metadata_col = struct_array + .column_by_name("metadata") + .ok_or_else(|| DataFusionError::Execution("Variant missing metadata column".to_string()))?; + let value_col = struct_array + .column_by_name("value") + .ok_or_else(|| DataFusionError::Execution("Variant missing value column".to_string()))?; + + let metadata_binary = metadata_col + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Execution("Variant metadata not BinaryView".to_string()))?; + let value_binary = value_col + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Execution("Variant value not BinaryView".to_string()))?; + + let mut builder = BooleanArray::builder(struct_array.len()); + + for i in 0..struct_array.len() { + if struct_array.is_null(i) { + builder.append_null(); + continue; + } + + let metadata = metadata_binary.value(i); + let value = value_binary.value(i); + + // Decode Variant to JSON + let variant = Variant::new(metadata, value); + let json_value = variant_to_serde_json(&variant); + + // Apply JSONPath and check if any matches exist + let matches = json_path.query(&json_value); + builder.append_value(!matches.is_empty()); + } + + Ok(Arc::new(builder.finish())) +} + +/// Evaluate JSONPath on a JSON string array +fn evaluate_jsonpath_on_json_string(array: &ArrayRef, json_path: &serde_json_path::JsonPath) -> datafusion::error::Result { + let mut builder = BooleanArray::builder(array.len()); + + // Handle different string types + if let Some(string_array) = array.as_any().downcast_ref::() { + for i in 0..string_array.len() { + if string_array.is_null(i) { + builder.append_null(); + } else { + let json_str = string_array.value(i); + let result = match serde_json::from_str::(json_str) { + Ok(json_value) => !json_path.query(&json_value).is_empty(), + Err(_) => false, // Invalid JSON returns false + }; + builder.append_value(result); + } + } + } else if let Some(string_array) = array.as_any().downcast_ref::() { + for i in 0..string_array.len() { + if string_array.is_null(i) { + builder.append_null(); + } else { + let json_str = string_array.value(i); + let result = match serde_json::from_str::(json_str) { + Ok(json_value) => !json_path.query(&json_value).is_empty(), + Err(_) => false, + }; + builder.append_value(result); + } + } + } else { + return Err(DataFusionError::Execution( + "jsonb_path_exists requires JSON string or Variant input".to_string(), + )); + } + + Ok(Arc::new(builder.finish())) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/schema_loader.rs b/src/schema_loader.rs index 088bd140..9a051fe2 100644 --- a/src/schema_loader.rs +++ b/src/schema_loader.rs @@ -104,6 +104,15 @@ fn parse_arrow_data_type(s: &str) -> anyhow::Result { "List(Utf8)" => ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Utf8View, true))), "Timestamp(Microsecond, None)" => ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, None), "Timestamp(Microsecond, Some(\"UTC\"))" => ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())), + // Variant Binary Encoding: Struct with metadata and value binary fields + // Using BinaryView for compatibility with datafusion-variant/parquet-variant-compute + "Variant" => ArrowDataType::Struct( + vec![ + Arc::new(Field::new("metadata", ArrowDataType::BinaryView, false)), + Arc::new(Field::new("value", ArrowDataType::BinaryView, false)), + ] + .into(), + ), _ => anyhow::bail!("Unknown type: {}", s), }) } @@ -116,6 +125,7 @@ fn parse_delta_data_type(s: &str) -> anyhow::Result { "Int32" | "UInt32" => DeltaDataType::Primitive(Integer), "Int64" | "UInt64" => DeltaDataType::Primitive(Long), "List(Utf8)" => DeltaDataType::Array(Box::new(ArrayType::new(DeltaDataType::Primitive(String), true))), + "Variant" => DeltaDataType::unshredded_variant(), _ if s.starts_with("Timestamp") => DeltaDataType::Primitive(Timestamp), _ => anyhow::bail!("Unknown type: {}", s), }) @@ -180,3 +190,19 @@ pub fn get_schema(table_name: &str) -> Option<&'static TableSchema> { pub fn get_default_schema() -> &'static TableSchema { registry().get_default().expect("No schemas available in registry") } + +/// Returns true if the given Arrow DataType represents a Variant type (Struct with metadata + value BinaryView fields) +pub fn is_variant_type(data_type: &ArrowDataType) -> bool { + match data_type { + ArrowDataType::Struct(fields) if fields.len() == 2 => { + fields.iter().any(|f| f.name() == "metadata" && matches!(f.data_type(), ArrowDataType::BinaryView)) + && fields.iter().any(|f| f.name() == "value" && matches!(f.data_type(), ArrowDataType::BinaryView)) + } + _ => false, + } +} + +/// Get indices of Variant columns in a schema +pub fn get_variant_column_indices(schema: &SchemaRef) -> Vec { + schema.fields().iter().enumerate().filter(|(_, f)| is_variant_type(f.data_type())).map(|(i, _)| i).collect() +} diff --git a/tests/connection_pressure_test.rs b/tests/connection_pressure_test.rs index 37e644b4..021804ad 100644 --- a/tests/connection_pressure_test.rs +++ b/tests/connection_pressure_test.rs @@ -47,7 +47,10 @@ mod connection_pressure { db.setup_session_context(&mut ctx).expect("Failed to setup context"); let opts = ServerOptions::new().with_port(port).with_host("0.0.0.0".to_string()); - let auth_config = timefusion::pgwire_handlers::AuthConfig::default(); + let auth_config = timefusion::pgwire_handlers::AuthConfig { + username: "postgres".into(), + password: Some("postgres".into()), + }; tokio::select! { _ = shutdown_clone.notified() => {}, diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 98153d4c..a161bfa6 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -65,7 +65,10 @@ mod integration { db_clone.setup_session_context(&mut ctx).expect("Failed to setup context"); let opts = ServerOptions::new().with_port(port).with_host("0.0.0.0".to_string()); - let auth_config = timefusion::pgwire_handlers::AuthConfig::default(); + let auth_config = timefusion::pgwire_handlers::AuthConfig { + username: "postgres".into(), + password: Some("postgres".into()), + }; tokio::select! { _ = shutdown_clone.notified() => {}, diff --git a/tests/slt/custom_functions.slt b/tests/slt/custom_functions.slt index 58d48eca..cbc734a4 100644 --- a/tests/slt/custom_functions.slt +++ b/tests/slt/custom_functions.slt @@ -8,10 +8,10 @@ INSERT INTO otel_logs_and_spans ( project_id, timestamp, id, hashes, date, parent_id, name, kind, resource___service___name, status_code, status_message, level, duration, summary -) VALUES +) VALUES ('test_functions', TIMESTAMP '2024-01-15T14:30:45.123456Z', 'func_test_1', ARRAY['hash1']::VARCHAR[], DATE '2024-01-15', NULL, 'test_to_char', 'SERVER', 'test-service', - 'OK', ARRAY['Test record'], 'INFO', 1000000, ARRAY['Test to_char function']), + 'OK', 'Test record', 'INFO', 1000000, ARRAY['Test to_char function']), ('test_functions', TIMESTAMP '2024-12-25T08:00:00Z', 'func_test_2', ARRAY['hash2']::VARCHAR[], DATE '2024-12-25', NULL, 'test_christmas', 'SERVER', 'test-service', 'OK', 'Christmas test', 'INFO', 2000000, ARRAY['Test date formatting']) @@ -120,10 +120,10 @@ INSERT INTO otel_logs_and_spans ( project_id, timestamp, id, hashes, date, parent_id, name, kind, resource___service___name, status_code, status_message, level, duration, summary -) VALUES +) VALUES ('test_formats', TIMESTAMP '2024-07-04T16:45:30Z', 'format_test_1', ARRAY['hash_format']::VARCHAR[], DATE '2024-07-04', NULL, 'test_formats', 'SERVER', 'format-service', - 'OK', ARRAY['Test various formats'], 'INFO', 1000000, ARRAY['Test different date formats']) + 'OK', 'Test various formats', 'INFO', 1000000, ARRAY['Test different date formats']) # Test various date format patterns query T diff --git a/tests/slt/function_availability_test.slt b/tests/slt/function_availability_test.slt index a968f949..5e21f151 100644 --- a/tests/slt/function_availability_test.slt +++ b/tests/slt/function_availability_test.slt @@ -6,10 +6,10 @@ INSERT INTO otel_logs_and_spans ( project_id, timestamp, id, hashes, date, parent_id, name, kind, resource___service___name, status_code, status_message, level, duration, summary -) VALUES +) VALUES ('test_funcs', TIMESTAMP '2024-01-15T14:30:45.123456Z', 'func_test_1', ARRAY['hash1']::VARCHAR[], DATE '2024-01-15', NULL, 'test_json', 'SERVER', 'test-service', - 'OK', ARRAY['Test message'], 'INFO', 1000000, ARRAY['Test functions']) + 'OK', 'Test message', 'INFO', 1000000, ARRAY['Test functions']) # === Test EXTRACT function (should work - DataFusion built-in) === diff --git a/tests/slt/integration.slt b/tests/slt/integration.slt index 2108fdd9..8d17f927 100644 --- a/tests/slt/integration.slt +++ b/tests/slt/integration.slt @@ -57,7 +57,7 @@ INSERT INTO otel_logs_and_spans ( name, resource___service___name, level, status_message, summary ) VALUES ( 'prod_monitoring', TIMESTAMP '2023-01-01T10:10:00Z', 'log_1', ARRAY[]::VARCHAR[], DATE '2023-01-01', - 'application.startup', 'user-service', 'INFO', ARRAY['Service started successfully'], ARRAY['Application startup log - INFO level'] + 'application.startup', 'user-service', 'INFO', 'Service started successfully', ARRAY['Application startup log - INFO level'] ) statement ok diff --git a/tests/slt/json_functions.slt b/tests/slt/json_functions.slt index 03b2339a..76581e91 100644 --- a/tests/slt/json_functions.slt +++ b/tests/slt/json_functions.slt @@ -57,7 +57,7 @@ INSERT INTO otel_logs_and_spans ( 'test_service', 'parent123', '2025-08-07T10:00:00Z', - '[{"event_name": "start"}, {"event_name": "exception"}]', + json_to_variant('[{"event_name": "start"}, {"event_name": "exception"}]'), ARRAY['{"status": "ok", "count": 5}'], 'span123', '00000000-0000-0000-0000-000000000001', @@ -91,7 +91,7 @@ INSERT INTO otel_logs_and_spans ( 'test_service2', 'parent456', '2025-08-07T11:00:00Z', - '[{"event_name": "info"}]', + json_to_variant('[{"event_name": "info"}]'), ARRAY['{"status": "error", "count": 0}'], 'span456', '00000000-0000-0000-0000-000000000002', diff --git a/tests/slt/variant_functions.slt b/tests/slt/variant_functions.slt new file mode 100644 index 00000000..efadf62e --- /dev/null +++ b/tests/slt/variant_functions.slt @@ -0,0 +1,269 @@ +# Test Variant type functions in TimeFusion +# This tests the Parquet Variant binary encoding and extraction functions + +# === Basic Variant Functions === + +# Test json_to_variant / variant_to_json round-trip with simple object +query T +SELECT variant_to_json(json_to_variant('{"key": "value"}')); +---- +{"key":"value"} + +# Test round-trip with nested object +query T +SELECT variant_to_json(json_to_variant('{"user": {"name": "Alice", "age": 30}}')); +---- +{"user":{"age":30,"name":"Alice"}} + +# Test round-trip with array +query T +SELECT variant_to_json(json_to_variant('[1, 2, 3, "four"]')); +---- +[1,2,3,"four"] + +# Test round-trip with primitives +query T +SELECT variant_to_json(json_to_variant('123')); +---- +123 + +query T +SELECT variant_to_json(json_to_variant('"hello"')); +---- +"hello" + +query T +SELECT variant_to_json(json_to_variant('true')); +---- +true + +query T +SELECT variant_to_json(json_to_variant('null')); +---- +null + +# === variant_get with path extraction === + +# Simple field extraction +query T +SELECT variant_to_json(variant_get(json_to_variant('{"name": "test", "value": 42}'), 'name')); +---- +"test" + +# Nested path extraction +query T +SELECT variant_to_json(variant_get(json_to_variant('{"a": {"b": {"c": "deep"}}}'), 'a.b.c')); +---- +"deep" + +# Array index extraction +query T +SELECT variant_to_json(variant_get(json_to_variant('{"items": [10, 20, 30]}'), 'items[0]')); +---- +10 + +query T +SELECT variant_to_json(variant_get(json_to_variant('{"items": [10, 20, 30]}'), 'items[2]')); +---- +30 + +# Non-existent path returns JSON null (via variant_to_json) +query T +SELECT variant_to_json(variant_get(json_to_variant('{"a": 1}'), 'nonexistent')); +---- +null + +# === is_variant_null === + +query B +SELECT is_variant_null(json_to_variant('null')); +---- +true + +query B +SELECT is_variant_null(json_to_variant('{"a": 1}')); +---- +false + +query B +SELECT is_variant_null(json_to_variant('0')); +---- +false + +query B +SELECT is_variant_null(json_to_variant('""')); +---- +false + +# === variant_pretty for debugging === + +query T +SELECT variant_pretty(json_to_variant('123')); +---- +Int8(123) + +# === jsonb_path_exists with literal variant === + +# Basic JSONPath query - check if path exists +query B +SELECT jsonb_path_exists(json_to_variant('{"user": {"name": "Alice"}}'), '$.user.name'); +---- +true + +# Check for non-existent path +query B +SELECT jsonb_path_exists(json_to_variant('{"user": {"name": "Alice"}}'), '$.nonexistent'); +---- +false + +# Array wildcard query - check if any item has 'name' field +query B +SELECT jsonb_path_exists(json_to_variant('{"items": [{"name": "a"}, {"name": "b"}]}'), '$.items[*].name'); +---- +true + +# Array wildcard with value check (RFC 9535 style) +query B +SELECT jsonb_path_exists(json_to_variant('[1, 2, 3]'), '$[*]'); +---- +true + +# JSONPath on null variant +query B +SELECT jsonb_path_exists(json_to_variant('null'), '$.any'); +---- +false + +# JSONPath on JSON string (not Variant) +query B +SELECT jsonb_path_exists('{"a": 1}', '$.a'); +---- +true + +# JSONPath on JSON string - non-existent +query B +SELECT jsonb_path_exists('{"a": 1}', '$.b'); +---- +false + +# === PostgreSQL-style Arrow Operators on Variant === + +# Test -> operator (returns Variant) +query T +SELECT variant_to_json(json_to_variant('{"user": {"name": "Alice", "id": 123}}')->'user'); +---- +{"id":123,"name":"Alice"} + +# Test chained -> operators (should be flattened to single variant_get) +query T +SELECT variant_to_json(json_to_variant('{"user": {"name": "Alice", "id": 123}}')->'user'->'name'); +---- +"Alice" + +# Test ->> operator (returns text via variant_to_json) +query T +SELECT json_to_variant('{"user": {"name": "Alice", "id": 123}}')->'user'->>'name'; +---- +"Alice" + +# Test array index access with -> operator +query T +SELECT variant_to_json(json_to_variant('{"items": [{"name": "item1", "qty": 5}, {"name": "item2", "qty": 10}]}')->'items'->0); +---- +{"name":"item1","qty":5} + +# Test chained array access with ->> for text +query T +SELECT json_to_variant('{"items": [{"name": "item1"}, {"name": "item2"}]}')->'items'->0->>'name'; +---- +"item1" + +# Test accessing second array element +query T +SELECT json_to_variant('{"items": [{"qty": 5}, {"qty": 10}]}')->'items'->1->>'qty'; +---- +10 + +# Test deep nesting with arrow operators +query T +SELECT json_to_variant('{"a": {"b": {"c": {"d": "deep"}}}}')->'a'->'b'->'c'->>'d'; +---- +"deep" + +# Test numeric extraction via ->> +query T +SELECT json_to_variant('{"count": 42}')->>'count'; +---- +42 + +# Test boolean extraction via ->> +query T +SELECT json_to_variant('{"active": true}')->>'active'; +---- +true + +# Test -> on non-existent field returns null variant +query T +SELECT variant_to_json(json_to_variant('{"a": 1}')->'nonexistent'); +---- +null + +# Test array with mixed types +query T +SELECT variant_to_json(json_to_variant('[1, "two", true, null]')->0); +---- +1 + +query T +SELECT json_to_variant('[1, "two", true, null]')->1->>''; +---- +"two" + +# Test complex nested structure +query T +SELECT json_to_variant('{"users": [{"profile": {"email": "alice@example.com"}}]}')->'users'->0->'profile'->>'email'; +---- +"alice@example.com" + +# Test -> followed by variant_get (mixed usage) +query T +SELECT variant_to_json(variant_get(json_to_variant('{"data": {"nested": {"value": 123}}}')->'data', 'nested.value')); +---- +123 + +# === Regex on extracted text (DataFusion native ~* operator) === + +query B +SELECT json_to_variant('{"name": "Alice"}')->>'name' ~* 'ali.*'; +---- +true + +query B +SELECT json_to_variant('{"message": "Error: Connection timeout"}')->>'message' ~* 'error.*timeout'; +---- +true + +query B +SELECT json_to_variant('{"message": "Success"}')->>'message' ~* 'error'; +---- +false + +# === Arrow operators on JSON strings (via datafusion-functions-json) === + +# Arrow operators on JSON strings +query T +SELECT '{"user": {"name": "Eve"}}'->'user'->>'name'; +---- +Eve + +# Chained arrows on JSON string +query T +SELECT '{"a": {"b": {"c": "deep"}}}'->'a'->'b'->>'c'; +---- +deep + +# Array access on JSON string with ->> +query T +SELECT '{"items": [10, 20, 30]}'->>'items'; +---- +[10, 20, 30] diff --git a/tests/sqllogictest.rs b/tests/sqllogictest.rs index 0cb07bc6..144df029 100644 --- a/tests/sqllogictest.rs +++ b/tests/sqllogictest.rs @@ -197,7 +197,10 @@ mod sqllogictest_tests { db.setup_session_context(&mut session_context).expect("Failed to setup session context"); let opts = ServerOptions::new().with_port(port).with_host("0.0.0.0".to_string()); - let auth_config = timefusion::pgwire_handlers::AuthConfig::default(); + let auth_config = timefusion::pgwire_handlers::AuthConfig { + username: "postgres".into(), + password: Some("postgres".into()), + }; // Wait for shutdown signal or server termination tokio::select! { From cfe89e735829460b2eaabc642bc648c5f1a78e4b Mon Sep 17 00:00:00 2001 From: Anthony Alaribe Date: Thu, 29 Jan 2026 23:24:39 -0800 Subject: [PATCH 2/7] Fix dependencies and add schema mismatch warning - Replace local path dependencies with git dependencies: - deltalake: Use fork with VariantType support (tonyalaribe/delta-rs) - datafusion-variant: Use git dependency with specific rev - Add warning log when schema has more fields than batch columns (helps debug schema evolution issues) --- Cargo.lock | 5 +++++ Cargo.toml | 6 +++--- src/database.rs | 4 ++++ 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1c01d908..77369b9e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2521,6 +2521,7 @@ dependencies = [ [[package]] name = "datafusion-variant" version = "0.1.0" +source = "git+https://github.com/tonyalaribe/datafusion-variant.git?rev=8b6b270#8b6b270f0f45693f6ccf39115d12bec9e9626012" dependencies = [ "arrow", "arrow-schema", @@ -2586,6 +2587,7 @@ dependencies = [ [[package]] name = "deltalake" version = "0.30.1" +source = "git+https://github.com/tonyalaribe/delta-rs.git?rev=ba769136c5dd9b84a7335ea67e42b67884bfcce3#ba769136c5dd9b84a7335ea67e42b67884bfcce3" dependencies = [ "ctor", "delta_kernel", @@ -2596,6 +2598,7 @@ dependencies = [ [[package]] name = "deltalake-aws" version = "0.13.0" +source = "git+https://github.com/tonyalaribe/delta-rs.git?rev=ba769136c5dd9b84a7335ea67e42b67884bfcce3#ba769136c5dd9b84a7335ea67e42b67884bfcce3" dependencies = [ "async-trait", "aws-config", @@ -2621,6 +2624,7 @@ dependencies = [ [[package]] name = "deltalake-core" version = "0.30.1" +source = "git+https://github.com/tonyalaribe/delta-rs.git?rev=ba769136c5dd9b84a7335ea67e42b67884bfcce3#ba769136c5dd9b84a7335ea67e42b67884bfcce3" dependencies = [ "arrow", "arrow-arith", @@ -2673,6 +2677,7 @@ dependencies = [ [[package]] name = "deltalake-derive" version = "0.30.0" +source = "git+https://github.com/tonyalaribe/delta-rs.git?rev=ba769136c5dd9b84a7335ea67e42b67884bfcce3#ba769136c5dd9b84a7335ea67e42b67884bfcce3" dependencies = [ "convert_case", "itertools 0.14.0", diff --git a/Cargo.toml b/Cargo.toml index a7596e4b..ffa18ab1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 Variant type support (adds variantType feature detection) -deltalake = { path = "/Users/tonyalaribe/Projects/apitoolkit/datafusion-projects/delta-rs/crates/deltalake", 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,7 +75,7 @@ bincode = { version = "2.0", features = ["serde"] } walrus-rust = "0.2.0" thiserror = "2.0" strum = { version = "0.27", features = ["derive"] } -datafusion-variant = { path = "/Users/tonyalaribe/Projects/apitoolkit/datafusion-projects/datafusion-variant" } +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" diff --git a/src/database.rs b/src/database.rs index b61f620c..1c20b645 100644 --- a/src/database.rs +++ b/src/database.rs @@ -97,6 +97,10 @@ pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> continue; } if idx >= columns.len() { + warn!( + "Schema mismatch: target schema has field '{}' at index {} but batch only has {} columns", + target_field.name(), idx, columns.len() + ); continue; } From 46a626de405a42c5a4c500bbf1f51a2de3d52a4c Mon Sep 17 00:00:00 2001 From: Anthony Alaribe Date: Thu, 29 Jan 2026 23:35:55 -0800 Subject: [PATCH 3/7] Apply rustfmt and fix clippy warning --- src/database.rs | 11 ++++++++--- src/functions.rs | 29 +++++++++++++++-------------- src/pgwire_handlers.rs | 25 ++++++++++++++++--------- tests/buffer_consistency_test.rs | 2 +- 4 files changed, 40 insertions(+), 27 deletions(-) diff --git a/src/database.rs b/src/database.rs index 1c20b645..f98b9b65 100644 --- a/src/database.rs +++ b/src/database.rs @@ -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; @@ -99,7 +99,9 @@ pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> if idx >= columns.len() { warn!( "Schema mismatch: target schema has field '{}' at index {} but batch only has {} columns", - target_field.name(), idx, columns.len() + target_field.name(), + idx, + columns.len() ); continue; } @@ -1283,7 +1285,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) { diff --git a/src/functions.rs b/src/functions.rs index 60eef686..e0a8e5ed 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -9,8 +9,8 @@ use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::common::{DFSchema, DataFusionError, ExprSchema, ScalarValue, not_impl_err}; use datafusion::logical_expr::ExprSchemable; use datafusion::logical_expr::{ - Accumulator, AggregateUDF, ColumnarValue, Expr, ScalarFunctionArgs, ScalarFunctionImplementation, - ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, Volatility, create_udaf, create_udf, + Accumulator, AggregateUDF, ColumnarValue, Expr, ScalarFunctionArgs, ScalarFunctionImplementation, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, + Volatility, create_udaf, create_udf, expr::{Alias, ScalarFunction}, planner::{ExprPlanner, PlannerResult, RawBinaryExpr}, }; @@ -121,23 +121,25 @@ fn extract_path_component(expr: &Expr) -> Option { fn is_variant_column(expr: &Expr, schema: &DFSchema) -> bool { match expr { // Direct column reference - Expr::Column(col) => schema - .field_from_column(col) - .map(|f| is_variant_type(f.data_type())) - .unwrap_or(false), + Expr::Column(col) => schema.field_from_column(col).map(|f| is_variant_type(f.data_type())).unwrap_or(false), // Unwrap aliases Expr::Alias(alias) => is_variant_column(&alias.expr, schema), // Check if it's a call to a variant-producing function Expr::ScalarFunction(func) => { let name = func.func.name(); - matches!(name, "json_to_variant" | "variant_get" | "cast_to_variant" - | "variant_object_construct" | "variant_list_construct" - | "variant_object_insert" | "variant_list_insert") + matches!( + name, + "json_to_variant" + | "variant_get" + | "cast_to_variant" + | "variant_object_construct" + | "variant_list_construct" + | "variant_object_insert" + | "variant_list_insert" + ) } // Try to get the type for other expressions - _ => expr.get_type(schema) - .map(|dt| is_variant_type(&dt)) - .unwrap_or(false), + _ => expr.get_type(schema).map(|dt| is_variant_type(&dt)).unwrap_or(false), } } @@ -1272,8 +1274,7 @@ impl ScalarUDFImpl for JsonbPathExistsUDF { }; // Parse the JSONPath expression - let json_path = serde_json_path::JsonPath::parse(&path_str) - .map_err(|e| DataFusionError::Execution(format!("Invalid JSONPath: {}", e)))?; + let json_path = serde_json_path::JsonPath::parse(&path_str).map_err(|e| DataFusionError::Execution(format!("Invalid JSONPath: {}", e)))?; // Process based on input type let result = if is_variant_type(json_array.data_type()) { diff --git a/src/pgwire_handlers.rs b/src/pgwire_handlers.rs index d9efa614..a578431e 100644 --- a/src/pgwire_handlers.rs +++ b/src/pgwire_handlers.rs @@ -1,5 +1,6 @@ use async_trait::async_trait; use datafusion::execution::context::SessionContext; +use datafusion_postgres::DfSessionService; use datafusion_postgres::pgwire::api::auth::cleartext::CleartextPasswordAuthStartupHandler; use datafusion_postgres::pgwire::api::auth::{AuthSource, DefaultServerParameterProvider, LoginInfo, Password, StartupHandler}; use datafusion_postgres::pgwire::api::portal::Portal; @@ -10,12 +11,11 @@ use datafusion_postgres::pgwire::api::store::PortalStore; use datafusion_postgres::pgwire::api::{ClientInfo, ClientPortalStore, ErrorHandler, PgWireServerHandlers}; use datafusion_postgres::pgwire::error::{PgWireError, PgWireResult}; use datafusion_postgres::pgwire::messages::PgWireBackendMessage; -use datafusion_postgres::DfSessionService; use futures::Sink; use std::fmt::Debug; use std::sync::Arc; use tracing::field::Empty; -use tracing::{info, instrument, Instrument}; +use tracing::{Instrument, info, instrument}; /// Auth configuration for PgWire server #[derive(Debug, Clone)] @@ -26,7 +26,10 @@ pub struct AuthConfig { impl Default for AuthConfig { fn default() -> Self { - Self { username: "postgres".into(), password: None } + Self { + username: "postgres".into(), + password: None, + } } } @@ -110,7 +113,9 @@ pub struct LoggingSimpleQueryHandler { impl LoggingSimpleQueryHandler { pub fn new(session_context: Arc) -> Self { - Self { inner: DfSessionService::new(session_context) } + Self { + inner: DfSessionService::new(session_context), + } } } @@ -176,7 +181,9 @@ pub struct LoggingExtendedQueryHandler { impl LoggingExtendedQueryHandler { pub fn new(session_context: Arc) -> Self { - Self { inner: DfSessionService::new(session_context) } + Self { + inner: DfSessionService::new(session_context), + } } } @@ -230,15 +237,15 @@ impl ExtendedQueryHandler for LoggingExtendedQueryHandler { span.record("query.text", sanitize_query(query, operation).as_str()); let execute_span = tracing::trace_span!(parent: &span, "datafusion.execute"); - ::do_query(&self.inner, client, portal, max_rows).instrument(execute_span).await + ::do_query(&self.inner, client, portal, max_rows) + .instrument(execute_span) + .await } } /// Start the server with custom handlers pub async fn serve_with_logging( - session_context: Arc, - options: &datafusion_postgres::ServerOptions, - auth_config: AuthConfig, + session_context: Arc, options: &datafusion_postgres::ServerOptions, auth_config: AuthConfig, ) -> Result<(), Box> { let handlers = Arc::new(LoggingHandlerFactory::new(session_context, auth_config)); datafusion_postgres::serve_with_handlers(handlers, options).await?; diff --git a/tests/buffer_consistency_test.rs b/tests/buffer_consistency_test.rs index dbcc6055..4d4d056d 100644 --- a/tests/buffer_consistency_test.rs +++ b/tests/buffer_consistency_test.rs @@ -21,7 +21,7 @@ async fn setup_db_with_buffer(mode: BufferMode) -> Result<(Arc, Arc Date: Thu, 29 Jan 2026 23:52:55 -0800 Subject: [PATCH 4/7] Improve Variant error handling and reduce code duplication - Replace .unwrap() with proper error handling in convert_variant_columns - json_strings_to_variant now fails fast on invalid JSON instead of silently inserting NULL - Upgrade schema mismatch logging from warn to error level - Expand registration order comment for clarity - Add scalar_to_string() helper to DRY string extraction from ScalarValue --- src/database.rs | 46 +++++++++++++++++++++++++--------------------- src/functions.rs | 38 ++++++++++++++++---------------------- 2 files changed, 41 insertions(+), 43 deletions(-) diff --git a/src/database.rs b/src/database.rs index f98b9b65..5ebf6172 100644 --- a/src/database.rs +++ b/src/database.rs @@ -97,8 +97,8 @@ pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> continue; } if idx >= columns.len() { - warn!( - "Schema mismatch: target schema has field '{}' at index {} but batch only has {} columns", + error!( + "Schema mismatch: target expects '{}' at index {} but batch has only {} columns (possible schema evolution issue)", target_field.name(), idx, columns.len() @@ -112,16 +112,22 @@ pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> // Only convert if source is a string type and target is Variant let converted: Option = match col_type { DataType::Utf8View => { - let arr = col.as_any().downcast_ref::().unwrap(); - Some(Arc::new(json_strings_to_variant(arr.iter()))) + let arr = col.as_any().downcast_ref::().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::().unwrap(); - Some(Arc::new(json_strings_to_variant(arr.iter()))) + let arr = col.as_any().downcast_ref::().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::().unwrap(); - Some(Arc::new(json_strings_to_variant(arr.iter()))) + let arr = col.as_any().downcast_ref::().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 }; @@ -136,27 +142,25 @@ pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> 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 -fn json_strings_to_variant<'a>(iter: impl Iterator>) -> datafusion::arrow::array::StructArray { +/// 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>) -> DFResult { use parquet_variant_compute::VariantArrayBuilder; use parquet_variant_json::JsonToVariant; let items: Vec<_> = iter.collect(); let mut builder = VariantArrayBuilder::new(items.len()); - for item in items { + for (row_idx, item) in items.into_iter().enumerate() { match item { - Some(json_str) => { - if let Err(e) = builder.append_json(json_str) { - warn!("Failed to parse JSON '{}': {}, inserting as null", json_str, e); - builder.append_null(); - } - } + 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(), } } - builder.build().into() + Ok(builder.build().into()) } // Compression level for parquet files - kept for WriterProperties fallback @@ -790,11 +794,11 @@ impl Database { self.register_pg_settings_table(ctx)?; self.register_set_config_udf(ctx); - // Register custom PostgreSQL-compatible functions BEFORE JSON functions - // so VariantAwareExprPlanner gets first chance at -> and ->> operators + // 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 (includes JsonExprPlanner for -> and ->> on string columns) + // JSON functions (JsonExprPlanner for -> and ->> on string columns - must come after Variant handlers) self.register_json_functions(ctx); Ok(()) diff --git a/src/functions.rs b/src/functions.rs index e0a8e5ed..d14cb3ef 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -22,6 +22,14 @@ use tdigests::TDigest; use crate::schema_loader::is_variant_type; +/// Extract a String from any ScalarValue string type (Utf8, Utf8View, LargeUtf8) +fn scalar_to_string(scalar: &ScalarValue) -> Option { + match scalar { + ScalarValue::Utf8(Some(s)) | ScalarValue::Utf8View(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => Some(s.clone()), + _ => None, + } +} + // ============================================================================ // Variant-Aware Expression Planner // ============================================================================ @@ -286,12 +294,8 @@ impl ScalarUDFImpl for ToCharUDF { // Extract format string let format_str = match &args[1] { - ColumnarValue::Scalar(scalar) => match scalar { - ScalarValue::Utf8(Some(s)) => s.clone(), - ScalarValue::Utf8View(Some(s)) => s.clone(), - ScalarValue::LargeUtf8(Some(s)) => s.clone(), - _ => return Err(DataFusionError::Execution("Format string must be a UTF8 string".to_string())), - }, + ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) + .ok_or_else(|| DataFusionError::Execution("Format string must be a UTF8 string".to_string()))?, ColumnarValue::Array(arr) => { if let Some(str_arr) = arr.as_any().downcast_ref::() { if str_arr.len() == 1 && !str_arr.is_null(0) { @@ -429,12 +433,8 @@ impl ScalarUDFImpl for AtTimeZoneUDF { // Extract timezone string let tz_str = match &args[1] { - ColumnarValue::Scalar(scalar) => match scalar { - ScalarValue::Utf8(Some(s)) => s.clone(), - ScalarValue::Utf8View(Some(s)) => s.clone(), - ScalarValue::LargeUtf8(Some(s)) => s.clone(), - _ => return Err(DataFusionError::Execution("Timezone must be a UTF8 string".to_string())), - }, + ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) + .ok_or_else(|| DataFusionError::Execution("Timezone must be a UTF8 string".to_string()))?, ColumnarValue::Array(arr) => { if let Some(str_arr) = arr.as_any().downcast_ref::() { if str_arr.len() == 1 && !str_arr.is_null(0) { @@ -870,12 +870,8 @@ fn create_time_bucket_udf() -> ScalarUDF { // Extract interval string let interval_str = match &args[0] { - ColumnarValue::Scalar(scalar) => match scalar { - datafusion::scalar::ScalarValue::Utf8(Some(s)) - | datafusion::scalar::ScalarValue::Utf8View(Some(s)) - | datafusion::scalar::ScalarValue::LargeUtf8(Some(s)) => s.clone(), - _ => return Err(DataFusionError::Execution("Interval must be a UTF8 string".to_string())), - }, + ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) + .ok_or_else(|| DataFusionError::Execution("Interval must be a UTF8 string".to_string()))?, ColumnarValue::Array(_) => { return Err(DataFusionError::Execution("Interval must be a scalar value".to_string())); } @@ -1264,10 +1260,8 @@ impl ScalarUDFImpl for JsonbPathExistsUDF { }; let path_str = match &args.args[1] { - ColumnarValue::Scalar(scalar) => match scalar { - ScalarValue::Utf8(Some(s)) | ScalarValue::Utf8View(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => s.clone(), - _ => return Err(DataFusionError::Execution("JSONPath must be a string".to_string())), - }, + ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) + .ok_or_else(|| DataFusionError::Execution("JSONPath must be a string".to_string()))?, ColumnarValue::Array(_) => { return Err(DataFusionError::Execution("JSONPath must be a scalar string".to_string())); } From 0aa753e11351145159e95b5b44060a28fa6797d4 Mon Sep 17 00:00:00 2001 From: Anthony Alaribe Date: Thu, 29 Jan 2026 23:53:57 -0800 Subject: [PATCH 5/7] Apply rustfmt --- src/database.rs | 50 +++++++++++++++++++++++++----------------------- src/functions.rs | 18 +++++++++-------- 2 files changed, 36 insertions(+), 32 deletions(-) diff --git a/src/database.rs b/src/database.rs index 5ebf6172..d4c60f99 100644 --- a/src/database.rs +++ b/src/database.rs @@ -110,27 +110,29 @@ pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> let col_type = col.data_type(); // Only convert if source is a string type and target is Variant - let converted: Option = match col_type { - DataType::Utf8View => { - let arr = col.as_any().downcast_ref::().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::().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::().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 - }; + let converted: Option = + match col_type { + DataType::Utf8View => { + let arr = col.as_any().downcast_ref::().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::() + .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::().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; @@ -153,9 +155,9 @@ fn json_strings_to_variant<'a>(iter: impl Iterator>) -> D 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)) - })?, + 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(), } } diff --git a/src/functions.rs b/src/functions.rs index d14cb3ef..1c8698bb 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -294,8 +294,9 @@ impl ScalarUDFImpl for ToCharUDF { // Extract format string let format_str = match &args[1] { - ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) - .ok_or_else(|| DataFusionError::Execution("Format string must be a UTF8 string".to_string()))?, + ColumnarValue::Scalar(scalar) => { + scalar_to_string(scalar).ok_or_else(|| DataFusionError::Execution("Format string must be a UTF8 string".to_string()))? + } ColumnarValue::Array(arr) => { if let Some(str_arr) = arr.as_any().downcast_ref::() { if str_arr.len() == 1 && !str_arr.is_null(0) { @@ -433,8 +434,9 @@ impl ScalarUDFImpl for AtTimeZoneUDF { // Extract timezone string let tz_str = match &args[1] { - ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) - .ok_or_else(|| DataFusionError::Execution("Timezone must be a UTF8 string".to_string()))?, + ColumnarValue::Scalar(scalar) => { + scalar_to_string(scalar).ok_or_else(|| DataFusionError::Execution("Timezone must be a UTF8 string".to_string()))? + } ColumnarValue::Array(arr) => { if let Some(str_arr) = arr.as_any().downcast_ref::() { if str_arr.len() == 1 && !str_arr.is_null(0) { @@ -870,8 +872,9 @@ fn create_time_bucket_udf() -> ScalarUDF { // Extract interval string let interval_str = match &args[0] { - ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) - .ok_or_else(|| DataFusionError::Execution("Interval must be a UTF8 string".to_string()))?, + ColumnarValue::Scalar(scalar) => { + scalar_to_string(scalar).ok_or_else(|| DataFusionError::Execution("Interval must be a UTF8 string".to_string()))? + } ColumnarValue::Array(_) => { return Err(DataFusionError::Execution("Interval must be a scalar value".to_string())); } @@ -1260,8 +1263,7 @@ impl ScalarUDFImpl for JsonbPathExistsUDF { }; let path_str = match &args.args[1] { - ColumnarValue::Scalar(scalar) => scalar_to_string(scalar) - .ok_or_else(|| DataFusionError::Execution("JSONPath must be a string".to_string()))?, + ColumnarValue::Scalar(scalar) => scalar_to_string(scalar).ok_or_else(|| DataFusionError::Execution("JSONPath must be a string".to_string()))?, ColumnarValue::Array(_) => { return Err(DataFusionError::Execution("JSONPath must be a scalar string".to_string())); } From 6a90c869f37e69698ee30d4b71dee39ba2dbdfde Mon Sep 17 00:00:00 2001 From: Anthony Alaribe Date: Fri, 30 Jan 2026 00:25:26 -0800 Subject: [PATCH 6/7] Add depth limit to variant_to_serde_json to prevent stack overflow - Add MAX_VARIANT_DEPTH (100) limit to prevent JSON bomb attacks via deeply nested Variant data - Remove noisy error log for expected schema evolution case in convert_variant_columns - Simplify ShortString handling to use to_string() directly --- src/database.rs | 8 ++------ src/functions.rs | 25 +++++++++++++++---------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/database.rs b/src/database.rs index d4c60f99..3434c7cf 100644 --- a/src/database.rs +++ b/src/database.rs @@ -96,13 +96,9 @@ pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> 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() { - error!( - "Schema mismatch: target expects '{}' at index {} but batch has only {} columns (possible schema evolution issue)", - target_field.name(), - idx, - columns.len() - ); continue; } diff --git a/src/functions.rs b/src/functions.rs index 1c8698bb..12a050e3 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -1285,12 +1285,18 @@ impl ScalarUDFImpl for JsonbPathExistsUDF { } } -/// Convert parquet_variant::Variant to serde_json::Value -fn variant_to_serde_json(variant: &parquet_variant::Variant) -> JsonValue { +const MAX_VARIANT_DEPTH: usize = 100; + +/// Convert parquet_variant::Variant to serde_json::Value with depth limit to prevent stack overflow +fn variant_to_serde_json(variant: &parquet_variant::Variant, depth: usize) -> Result { use base64::Engine; use parquet_variant::Variant; - match variant { + if depth > MAX_VARIANT_DEPTH { + return Err(DataFusionError::Execution(format!("Variant nesting depth exceeds limit of {}", MAX_VARIANT_DEPTH))); + } + + Ok(match variant { Variant::Null => JsonValue::Null, Variant::BooleanTrue => JsonValue::Bool(true), Variant::BooleanFalse => JsonValue::Bool(false), @@ -1312,19 +1318,19 @@ fn variant_to_serde_json(variant: &parquet_variant::Variant) -> JsonValue { Variant::TimestampNtzNanos(v) => json!(*v), Variant::Binary(bytes) => json!(base64::engine::general_purpose::STANDARD.encode(bytes)), Variant::String(s) => JsonValue::String(s.to_string()), - Variant::ShortString(s) => JsonValue::String(s.as_str().to_string()), + Variant::ShortString(s) => JsonValue::String(s.to_string()), Variant::Object(obj) => { let mut map = serde_json::Map::new(); for (key, value) in obj.iter() { - map.insert(key.to_string(), variant_to_serde_json(&value)); + map.insert(key.to_string(), variant_to_serde_json(&value, depth + 1)?); } JsonValue::Object(map) } Variant::List(list) => { - let items: Vec = list.iter().map(|v| variant_to_serde_json(&v)).collect(); + let items: Vec = list.iter().map(|v| variant_to_serde_json(&v, depth + 1)).collect::>()?; JsonValue::Array(items) } - } + }) } /// Evaluate JSONPath on a Variant (Struct) array @@ -1366,11 +1372,10 @@ fn evaluate_jsonpath_on_variant(array: &ArrayRef, json_path: &serde_json_path::J // Decode Variant to JSON let variant = Variant::new(metadata, value); - let json_value = variant_to_serde_json(&variant); + let json_value = variant_to_serde_json(&variant, 0)?; // Apply JSONPath and check if any matches exist - let matches = json_path.query(&json_value); - builder.append_value(!matches.is_empty()); + builder.append_value(!json_path.query(&json_value).is_empty()); } Ok(Arc::new(builder.finish())) From 2b72658a4832807a2bb696dd186177b54c99ce56 Mon Sep 17 00:00:00 2001 From: Anthony Alaribe Date: Fri, 30 Jan 2026 00:26:31 -0800 Subject: [PATCH 7/7] Apply rustfmt --- src/functions.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/functions.rs b/src/functions.rs index 12a050e3..d3bd8cc7 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -1293,7 +1293,10 @@ fn variant_to_serde_json(variant: &parquet_variant::Variant, depth: usize) -> Re use parquet_variant::Variant; if depth > MAX_VARIANT_DEPTH { - return Err(DataFusionError::Execution(format!("Variant nesting depth exceeds limit of {}", MAX_VARIANT_DEPTH))); + return Err(DataFusionError::Execution(format!( + "Variant nesting depth exceeds limit of {}", + MAX_VARIANT_DEPTH + ))); } Ok(match variant {