feat: variant type support extension types

Author: b41sh

Cargo.toml

@@ -638,6 +638,7 @@ color-eyre = { git = "https://github.com/eyre-rs/eyre.git", rev = "e5d92c3" }
 deltalake = { git = "https://github.com/delta-io/delta-rs", rev = "c149502" }
 display-more = { git = "https://github.com/databendlabs/display-more", tag = "v0.1.2" }
 ethnum = { git = "https://github.com/datafuse-extras/ethnum-rs", rev = "4cb05f1" }
+jsonb = { git = "https://github.com/b41sh/jsonb", rev = "4e9315efa67c46777fd3097f0d189ff5990ac24c" }
 map-api = { git = "https://github.com/databendlabs/map-api", tag = "v0.2.3" }
 openai_api_rust = { git = "https://github.com/datafuse-extras/openai-api", rev = "819a0ed" }
 openraft = { git = "https://github.com/databendlabs/openraft", tag = "v0.10.0-alpha.9" }

src/query/expression/src/types/variant.rs

@@ -230,13 +230,37 @@ pub fn cast_scalar_to_variant(scalar: ScalarRef, tz: &TimeZone, buf: &mut Vec<u8
             NumberScalar::Float32(n) => n.0.into(),
             NumberScalar::Float64(n) => n.0.into(),
         },
-        ScalarRef::Decimal(x) => x.to_float64().into(),
+        ScalarRef::Decimal(x) => {
+            match x {
+                DecimalScalar::Decimal128(value, size) => {
+                    let dec = jsonb::Decimal128 {
+                        scale: size.scale,
+                        value,
+                    };
+                    jsonb::Value::Number(jsonb::Number::Decimal128(dec))
+                }
+                DecimalScalar::Decimal256(value, size) => {
+                    let dec = jsonb::Decimal256 {
+                        scale: size.scale,
+                        value,
+                    };
+                    jsonb::Value::Number(jsonb::Number::Decimal256(dec))
+                }
+            }
+        }
         ScalarRef::Boolean(b) => jsonb::Value::Bool(b),
-        ScalarRef::Binary(s) => jsonb::Value::String(hex::encode_upper(s).into()),
+        ScalarRef::Binary(s) => jsonb::Value::Binary(s),
         ScalarRef::String(s) => jsonb::Value::String(s.into()),
-        ScalarRef::Timestamp(ts) => timestamp_to_string(ts, tz).to_string().into(),
-        ScalarRef::Date(d) => date_to_string(d, tz).to_string().into(),
-        ScalarRef::Interval(i) => interval_to_string(&i).to_string().into(),
+        ScalarRef::Timestamp(ts) => jsonb::Value::Timestamp(jsonb::Timestamp {offset:0, value: ts}),
+        ScalarRef::Date(d) => jsonb::Value::Timestamp(jsonb::Date {offset:0, value: d}),,
+        ScalarRef::Interval(i) => {
+            let interval = jsonb::Interval {
+                months: i.months(),
+                days: i.days(),
+                microseconds: i.microseconds(),
+            };
+            jsonb::Value::Interval(interval)
+        }
         ScalarRef::Array(col) => {
             let items = cast_scalars_to_variants(col.iter(), tz);
             let owned_jsonb = OwnedJsonb::build_array(items.iter().map(RawJsonb::new))

src/query/functions/src/scalars/variant.rs

@@ -729,6 +729,57 @@ pub fn register(registry: &mut FunctionRegistry) {
         }),
     );
 
+    registry.register_combine_nullable_1_arg::<VariantType, BinaryType, _, _>(
+        "as_binary",
+        |_, _| FunctionDomain::Full,
+        vectorize_with_builder_1_arg::<VariantType, NullableType<BinaryType>>(|v, output, ctx| {
+            if let Some(validity) = &ctx.validity {
+                if !validity.get_bit(output.len()) {
+                    output.push_null();
+                    return;
+                }
+            }
+            match RawJsonb::new(v).as_binary() {
+                Ok(Some(res)) => output.push(&res),
+                _ => output.push_null(),
+            }
+        }),
+    );
+
+    registry.register_combine_nullable_1_arg::<VariantType, DateType, _, _>(
+        "as_date",
+        |_, _| FunctionDomain::Full,
+        vectorize_with_builder_1_arg::<VariantType, NullableType<DateType>>(|v, output, ctx| {
+            if let Some(validity) = &ctx.validity {
+                if !validity.get_bit(output.len()) {
+                    output.push_null();
+                    return;
+                }
+            }
+            match RawJsonb::new(v).as_date() {
+                Ok(Some(res)) => output.push(res.value),
+                _ => output.push_null(),
+            }
+        }),
+    );
+
+    registry.register_combine_nullable_1_arg::<VariantType, TimestampType, _, _>(
+        "as_timestamp",
+        |_, _| FunctionDomain::Full,
+        vectorize_with_builder_1_arg::<VariantType, NullableType<TimestampType>>(|v, output, ctx| {
+            if let Some(validity) = &ctx.validity {
+                if !validity.get_bit(output.len()) {
+                    output.push_null();
+                    return;
+                }
+            }
+            match RawJsonb::new(v).as_timestamp() {
+                Ok(Some(res)) => output.push(res.value),
+                _ => output.push_null(),
+            }
+        }),
+    );
+
     registry.register_combine_nullable_1_arg::<VariantType, VariantType, _, _>(
         "as_array",
         |_, _| FunctionDomain::Full,