vortex_datafusion/convert/
scalars.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use arrow_array::Array;
7use arrow_array::StructArray;
8use arrow_schema::Field;
9use arrow_schema::Fields;
10use datafusion_common::ScalarValue;
11use vortex::buffer::ByteBuffer;
12use vortex::dtype::DType;
13use vortex::dtype::DecimalDType;
14use vortex::dtype::NativeDecimalType;
15use vortex::dtype::Nullability;
16use vortex::dtype::PType;
17use vortex::dtype::half::f16;
18use vortex::dtype::i256;
19use vortex::error::VortexExpect;
20use vortex::error::VortexResult;
21use vortex::error::vortex_bail;
22use vortex::error::vortex_err;
23use vortex::error::vortex_panic;
24use vortex::extension::datetime::AnyTemporal;
25use vortex::extension::datetime::TemporalMetadata;
26use vortex::extension::datetime::TimeUnit;
27use vortex::scalar::DecimalValue;
28use vortex::scalar::Scalar;
29use vortex_arrow::FromArrowType;
30
31use crate::convert::FromDataFusion;
32use crate::convert::TryToDataFusion;
33
34impl TryToDataFusion<ScalarValue> for Scalar {
35    fn try_to_df(&self) -> VortexResult<ScalarValue> {
36        Ok(match self.dtype() {
37            DType::Null => ScalarValue::Null,
38            DType::Bool(_) => ScalarValue::Boolean(self.as_bool().value()),
39            DType::Primitive(ptype, _) => {
40                let pscalar = self.as_primitive();
41                match ptype {
42                    PType::U8 => ScalarValue::UInt8(pscalar.typed_value::<u8>()),
43                    PType::U16 => ScalarValue::UInt16(pscalar.typed_value::<u16>()),
44                    PType::U32 => ScalarValue::UInt32(pscalar.typed_value::<u32>()),
45                    PType::U64 => ScalarValue::UInt64(pscalar.typed_value::<u64>()),
46                    PType::I8 => ScalarValue::Int8(pscalar.typed_value::<i8>()),
47                    PType::I16 => ScalarValue::Int16(pscalar.typed_value::<i16>()),
48                    PType::I32 => ScalarValue::Int32(pscalar.typed_value::<i32>()),
49                    PType::I64 => ScalarValue::Int64(pscalar.typed_value::<i64>()),
50                    PType::F16 => ScalarValue::Float16(pscalar.typed_value::<f16>()),
51                    PType::F32 => ScalarValue::Float32(pscalar.typed_value::<f32>()),
52                    PType::F64 => ScalarValue::Float64(pscalar.typed_value::<f64>()),
53                }
54            }
55            DType::Decimal(decimal_type, _) => {
56                let dscalar = self.as_decimal();
57                let precision = decimal_type.precision();
58                let scale = decimal_type.scale();
59
60                if precision <= i32::MAX_PRECISION {
61                    match dscalar.decimal_value() {
62                        None => ScalarValue::Decimal32(None, precision, scale),
63                        Some(value) => match value.cast::<i32>() {
64                            Some(v32) => ScalarValue::Decimal32(Some(v32), precision, scale),
65                            None => {
66                                vortex_bail!(
67                                    "invalid ScalarValue {value} for decimal with precision {precision}",
68                                )
69                            }
70                        },
71                    }
72                } else if precision <= i64::MAX_PRECISION {
73                    match dscalar.decimal_value() {
74                        None => ScalarValue::Decimal64(None, precision, scale),
75                        Some(value) => match value.cast::<i64>() {
76                            Some(v64) => ScalarValue::Decimal64(Some(v64), precision, scale),
77                            None => {
78                                vortex_bail!(
79                                    "invalid ScalarValue {value} for decimal with precision {precision}",
80                                )
81                            }
82                        },
83                    }
84                } else if precision <= i128::MAX_PRECISION {
85                    match dscalar.decimal_value() {
86                        None => ScalarValue::Decimal128(None, precision, scale),
87                        Some(value) => match value.cast::<i128>() {
88                            Some(v128) => ScalarValue::Decimal128(Some(v128), precision, scale),
89                            None => {
90                                vortex_bail!(
91                                    "invalid ScalarValue {value} for decimal with precision {precision}",
92                                )
93                            }
94                        },
95                    }
96                } else {
97                    match dscalar.decimal_value() {
98                        None => ScalarValue::Decimal256(None, precision, scale),
99                        Some(value) => match value.cast::<i256>() {
100                            Some(v256) => {
101                                ScalarValue::Decimal256(Some(v256.into()), precision, scale)
102                            }
103                            None => {
104                                vortex_bail!(
105                                    "invalid ScalarValue {value} for decimal with precision {precision}",
106                                )
107                            }
108                        },
109                    }
110                }
111            }
112            // SAFETY: By construction Utf8 scalar values are utf8
113            DType::Utf8(_) => ScalarValue::Utf8(self.as_utf8().value().cloned().map(|s| unsafe {
114                String::from_utf8_unchecked(Vec::<u8>::from(s.into_inner().into_inner()))
115            })),
116            DType::Binary(_) => ScalarValue::Binary(
117                self.as_binary()
118                    .value()
119                    .cloned()
120                    .map(|b| Vec::<u8>::from(b.into_inner())),
121            ),
122            dtype @ DType::List(..) => vortex_bail!(
123                "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type"
124            ),
125            dtype @ DType::FixedSizeList(..) => vortex_bail!(
126                "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type"
127            ),
128            DType::Struct(..) => struct_to_df(self)?,
129            dtype @ DType::Union(..) => vortex_bail!(
130                "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type"
131            ),
132            DType::Variant(_) => vortex_bail!("Variant scalars aren't supported with DF"),
133            DType::Extension(ext) => {
134                let storage_scalar = self.as_extension().to_storage_scalar();
135
136                let Some(temporal) = ext.metadata_opt::<AnyTemporal>() else {
137                    // Unknown extension type: perform scalar conversion using the canonical
138                    // scalar DType.
139                    return storage_scalar.try_to_df();
140                };
141
142                // Special handling: temporal extension types in Vortex correspond to Arrow's
143                // temporal physical types.
144                let pv = storage_scalar.as_primitive();
145                match temporal {
146                    TemporalMetadata::Timestamp(unit, tz) => match unit {
147                        TimeUnit::Nanoseconds => {
148                            ScalarValue::TimestampNanosecond(pv.as_::<i64>(), tz.clone())
149                        }
150                        TimeUnit::Microseconds => {
151                            ScalarValue::TimestampMicrosecond(pv.as_::<i64>(), tz.clone())
152                        }
153                        TimeUnit::Milliseconds => {
154                            ScalarValue::TimestampMillisecond(pv.as_::<i64>(), tz.clone())
155                        }
156                        TimeUnit::Seconds => {
157                            ScalarValue::TimestampSecond(pv.as_::<i64>(), tz.clone())
158                        }
159                        TimeUnit::Days => {
160                            unreachable!("Unsupported TimeUnit {unit} for {}", ext.id())
161                        }
162                    },
163                    TemporalMetadata::Date(unit) => match unit {
164                        TimeUnit::Milliseconds => ScalarValue::Date64(pv.as_::<i64>()),
165                        TimeUnit::Days => ScalarValue::Date32(pv.as_::<i32>()),
166                        _ => unreachable!("Unsupported TimeUnit {unit} for {}", ext.id()),
167                    },
168                    TemporalMetadata::Time(unit) => match unit {
169                        TimeUnit::Nanoseconds => ScalarValue::Time64Nanosecond(pv.as_::<i64>()),
170                        TimeUnit::Microseconds => ScalarValue::Time64Microsecond(pv.as_::<i64>()),
171                        TimeUnit::Milliseconds => ScalarValue::Time32Millisecond(pv.as_::<i32>()),
172                        TimeUnit::Seconds => ScalarValue::Time32Second(pv.as_::<i32>()),
173                        TimeUnit::Days => {
174                            unreachable!("Unsupported TimeUnit {unit} for {}", ext.id())
175                        }
176                    },
177                }
178            }
179        })
180    }
181}
182
183impl FromDataFusion<ScalarValue> for Scalar {
184    fn from_df(value: &ScalarValue) -> Scalar {
185        match value {
186            ScalarValue::Null => Scalar::null(DType::Null),
187            ScalarValue::Boolean(b) => b
188                .map(Scalar::from)
189                .unwrap_or_else(|| Scalar::null(DType::Bool(Nullability::Nullable))),
190            ScalarValue::Float16(f) => f.map(Scalar::from).unwrap_or_else(|| {
191                Scalar::null(DType::Primitive(PType::F16, Nullability::Nullable))
192            }),
193            ScalarValue::Float32(f) => f.map(Scalar::from).unwrap_or_else(|| {
194                Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable))
195            }),
196            ScalarValue::Float64(f) => f.map(Scalar::from).unwrap_or_else(|| {
197                Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))
198            }),
199            ScalarValue::Int8(i) => i.map(Scalar::from).unwrap_or_else(|| {
200                Scalar::null(DType::Primitive(PType::I8, Nullability::Nullable))
201            }),
202            ScalarValue::Int16(i) => i.map(Scalar::from).unwrap_or_else(|| {
203                Scalar::null(DType::Primitive(PType::I16, Nullability::Nullable))
204            }),
205            ScalarValue::Int32(i) => i.map(Scalar::from).unwrap_or_else(|| {
206                Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))
207            }),
208            ScalarValue::Int64(i) => i.map(Scalar::from).unwrap_or_else(|| {
209                Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))
210            }),
211            ScalarValue::UInt8(i) => i.map(Scalar::from).unwrap_or_else(|| {
212                Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable))
213            }),
214            ScalarValue::UInt16(i) => i.map(Scalar::from).unwrap_or_else(|| {
215                Scalar::null(DType::Primitive(PType::U16, Nullability::Nullable))
216            }),
217            ScalarValue::UInt32(i) => i.map(Scalar::from).unwrap_or_else(|| {
218                Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable))
219            }),
220            ScalarValue::UInt64(i) => i.map(Scalar::from).unwrap_or_else(|| {
221                Scalar::null(DType::Primitive(PType::U64, Nullability::Nullable))
222            }),
223            ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) | ScalarValue::LargeUtf8(s) => s
224                .as_ref()
225                .map(|s| Scalar::from(s.as_str()))
226                .unwrap_or_else(|| Scalar::null(DType::Utf8(Nullability::Nullable))),
227            ScalarValue::Binary(b)
228            | ScalarValue::BinaryView(b)
229            | ScalarValue::LargeBinary(b)
230            | ScalarValue::FixedSizeBinary(_, b) => b
231                .as_ref()
232                .map(|b| Scalar::binary(ByteBuffer::from(b.clone()), Nullability::Nullable))
233                .unwrap_or_else(|| Scalar::null(DType::Binary(Nullability::Nullable))),
234            ScalarValue::Date32(v)
235            | ScalarValue::Time32Second(v)
236            | ScalarValue::Time32Millisecond(v) => {
237                let dtype = DType::from_arrow((&value.data_type(), Nullability::Nullable));
238                Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from))
239                    .vortex_expect("unable to create a time `Scalar`")
240            }
241            ScalarValue::Date64(v)
242            | ScalarValue::Time64Microsecond(v)
243            | ScalarValue::Time64Nanosecond(v)
244            | ScalarValue::TimestampSecond(v, _)
245            | ScalarValue::TimestampMillisecond(v, _)
246            | ScalarValue::TimestampMicrosecond(v, _)
247            | ScalarValue::TimestampNanosecond(v, _) => {
248                let dtype = DType::from_arrow((&value.data_type(), Nullability::Nullable));
249                Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from))
250                    .vortex_expect("unable to create a time `Scalar`")
251            }
252            ScalarValue::Decimal32(decimal, precision, scale) => {
253                let decimal_dtype = DecimalDType::new(*precision, *scale);
254                let nullable = Nullability::Nullable;
255                if let Some(value) = decimal {
256                    Scalar::decimal(
257                        DecimalValue::I32(*value),
258                        decimal_dtype,
259                        Nullability::Nullable,
260                    )
261                } else {
262                    Scalar::null(DType::Decimal(decimal_dtype, nullable))
263                }
264            }
265            ScalarValue::Decimal64(decimal, precision, scale) => {
266                let decimal_dtype = DecimalDType::new(*precision, *scale);
267                let nullable = Nullability::Nullable;
268                if let Some(value) = decimal {
269                    Scalar::decimal(
270                        DecimalValue::I64(*value),
271                        decimal_dtype,
272                        Nullability::Nullable,
273                    )
274                } else {
275                    Scalar::null(DType::Decimal(decimal_dtype, nullable))
276                }
277            }
278            ScalarValue::Decimal128(decimal, precision, scale) => {
279                let decimal_dtype = DecimalDType::new(*precision, *scale);
280                let nullable = Nullability::Nullable;
281                if let Some(value) = decimal {
282                    Scalar::decimal(
283                        DecimalValue::I128(*value),
284                        decimal_dtype,
285                        Nullability::Nullable,
286                    )
287                } else {
288                    Scalar::null(DType::Decimal(decimal_dtype, nullable))
289                }
290            }
291            ScalarValue::Decimal256(decimal, precision, scale) => {
292                let decimal_dtype = DecimalDType::new(*precision, *scale);
293                let nullable = Nullability::Nullable;
294                if let Some(value) = decimal {
295                    Scalar::decimal(
296                        DecimalValue::I256(i256::from_le_bytes(value.to_le_bytes())),
297                        decimal_dtype,
298                        Nullability::Nullable,
299                    )
300                } else {
301                    Scalar::null(DType::Decimal(decimal_dtype, nullable))
302                }
303            }
304            ScalarValue::Dictionary(_, v) => Scalar::from_df(v.as_ref()),
305            ScalarValue::Struct(array) => struct_from_df(array),
306            _ => unimplemented!("Can't convert {value:?} value to a Vortex scalar"),
307        }
308    }
309}
310
311/// Converts a Vortex struct scalar to a DataFusion `ScalarValue::Struct`.
312fn struct_to_df(scalar: &Scalar) -> VortexResult<ScalarValue> {
313    let scalar = scalar.as_struct();
314    let struct_fields = scalar.struct_fields();
315    let (fields, arrays): (Vec<Field>, Vec<_>) = struct_fields
316        .names()
317        .iter()
318        .zip(struct_fields.fields())
319        .enumerate()
320        .map(|(idx, (name, field_dtype))| {
321            let nullable = field_dtype.is_nullable();
322            let child = if scalar.is_null() {
323                Scalar::null(field_dtype)
324            } else {
325                scalar
326                    .field_by_idx(idx)
327                    .ok_or_else(|| vortex_err!("missing struct field {name}"))?
328            };
329            let array = child
330                .try_to_df()?
331                .to_array()
332                .map_err(|e| vortex_err!("failed to build struct field array: {e}"))?;
333            Ok((
334                Field::new(name.as_ref(), array.data_type().clone(), nullable),
335                array,
336            ))
337        })
338        .collect::<VortexResult<Vec<_>>>()?
339        .into_iter()
340        .unzip();
341
342    let fields = Fields::from(fields);
343    let struct_array = if scalar.is_null() {
344        StructArray::new_null(fields, 1)
345    } else {
346        StructArray::try_new(fields, arrays, None)
347            .map_err(|e| vortex_err!("failed to build struct scalar array: {e}"))?
348    };
349    Ok(ScalarValue::Struct(Arc::new(struct_array)))
350}
351
352/// Converts a DataFusion `ScalarValue::Struct` (a one-row struct array) to a Vortex struct scalar.
353fn struct_from_df(array: &StructArray) -> Scalar {
354    let dtype = DType::from_arrow((array.data_type(), Nullability::Nullable));
355    if array.is_null(0) {
356        Scalar::null(dtype)
357    } else {
358        let children = array
359            .columns()
360            .iter()
361            .map(|column| {
362                Scalar::from_df(
363                    &ScalarValue::try_from_array(column.as_ref(), 0).unwrap_or_else(|e| {
364                        vortex_panic!("cannot convert struct field to a Vortex scalar: {e}")
365                    }),
366                )
367            })
368            .collect::<Vec<_>>();
369        Scalar::struct_(dtype, children)
370    }
371}
372
373#[cfg(test)]
374mod tests {
375    use datafusion_common::ScalarValue;
376    use datafusion_common::arrow::datatypes::i256 as arrow_i256;
377    use rstest::rstest;
378    use vortex::buffer::ByteBuffer;
379    use vortex::dtype::DType;
380    use vortex::dtype::DecimalDType;
381    use vortex::dtype::FieldNames;
382    use vortex::dtype::Nullability;
383    use vortex::dtype::PType;
384    use vortex::dtype::StructFields;
385    use vortex::dtype::UnionVariants;
386    use vortex::dtype::i256;
387    use vortex::scalar::DecimalValue;
388    use vortex::scalar::Scalar;
389
390    use super::*;
391
392    #[rstest]
393    #[case::u8_some(Scalar::from(42u8), ScalarValue::UInt8(Some(42)))]
394    #[case::u8_null(
395        Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)),
396        ScalarValue::UInt8(None)
397    )]
398    #[case::u16_some(Scalar::from(1234u16), ScalarValue::UInt16(Some(1234)))]
399    #[case::u16_null(
400        Scalar::null(DType::Primitive(PType::U16, Nullability::Nullable)),
401        ScalarValue::UInt16(None)
402    )]
403    #[case::u32_some(Scalar::from(123456u32), ScalarValue::UInt32(Some(123456)))]
404    #[case::u32_null(
405        Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)),
406        ScalarValue::UInt32(None)
407    )]
408    #[case::u64_some(Scalar::from(12345678u64), ScalarValue::UInt64(Some(12345678)))]
409    #[case::u64_null(
410        Scalar::null(DType::Primitive(PType::U64, Nullability::Nullable)),
411        ScalarValue::UInt64(None)
412    )]
413    #[case::i8_some(Scalar::from(-42i8), ScalarValue::Int8(Some(-42)))]
414    #[case::i8_null(
415        Scalar::null(DType::Primitive(PType::I8, Nullability::Nullable)),
416        ScalarValue::Int8(None)
417    )]
418    #[case::i16_some(Scalar::from(-1234i16), ScalarValue::Int16(Some(-1234)))]
419    #[case::i16_null(
420        Scalar::null(DType::Primitive(PType::I16, Nullability::Nullable)),
421        ScalarValue::Int16(None)
422    )]
423    #[case::i32_some(Scalar::from(-123456i32), ScalarValue::Int32(Some(-123456)))]
424    #[case::i32_null(
425        Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
426        ScalarValue::Int32(None)
427    )]
428    #[case::i64_some(Scalar::from(-12345678i64), ScalarValue::Int64(Some(-12345678)))]
429    #[case::i64_null(
430        Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)),
431        ScalarValue::Int64(None)
432    )]
433    #[case::f32_some(Scalar::from(1.5f32), ScalarValue::Float32(Some(1.5)))]
434    #[case::f32_null(
435        Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable)),
436        ScalarValue::Float32(None)
437    )]
438    #[case::f64_some(Scalar::from(2.5f64), ScalarValue::Float64(Some(2.5)))]
439    #[case::f64_null(
440        Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)),
441        ScalarValue::Float64(None)
442    )]
443    fn test_primitive_to_datafusion(
444        #[case] vortex_scalar: Scalar,
445        #[case] expected_df_scalar: ScalarValue,
446    ) {
447        let result = vortex_scalar.try_to_df().unwrap();
448        assert_eq!(result, expected_df_scalar);
449    }
450
451    #[rstest]
452    #[case::bool_some(Scalar::from(true), ScalarValue::Boolean(Some(true)))]
453    #[case::bool_some_false(Scalar::from(false), ScalarValue::Boolean(Some(false)))]
454    #[case::bool_null(
455        Scalar::null(DType::Bool(Nullability::Nullable)),
456        ScalarValue::Boolean(None)
457    )]
458    #[case::null_type(Scalar::null(DType::Null), ScalarValue::Null)]
459    fn test_bool_and_null_to_datafusion(
460        #[case] vortex_scalar: Scalar,
461        #[case] expected_df_scalar: ScalarValue,
462    ) {
463        let result = vortex_scalar.try_to_df().unwrap();
464        assert_eq!(result, expected_df_scalar);
465    }
466
467    #[rstest]
468    #[case::utf8_some(Scalar::from("hello"), ScalarValue::Utf8(Some("hello".to_string())))]
469    #[case::utf8_null(
470        Scalar::null(DType::Utf8(Nullability::Nullable)),
471        ScalarValue::Utf8(None)
472    )]
473    #[case::binary_some(
474        Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3, 4]), Nullability::NonNullable),
475        ScalarValue::Binary(Some(vec![1u8, 2, 3, 4]))
476    )]
477    #[case::binary_null(
478        Scalar::null(DType::Binary(Nullability::Nullable)),
479        ScalarValue::Binary(None)
480    )]
481    fn test_string_and_binary_to_datafusion(
482        #[case] vortex_scalar: Scalar,
483        #[case] expected_df_scalar: ScalarValue,
484    ) {
485        let result = vortex_scalar.try_to_df().unwrap();
486        assert_eq!(result, expected_df_scalar);
487    }
488
489    #[rstest]
490    #[case::decimal32_some(
491        Scalar::decimal(
492            DecimalValue::I32(1234),
493            DecimalDType::new(5, 2),
494            Nullability::NonNullable
495        ),
496        ScalarValue::Decimal32(Some(1234), 5, 2)
497    )]
498    #[case::decimal32_null(
499        Scalar::null(DType::Decimal(DecimalDType::new(5, 2), Nullability::Nullable)),
500        ScalarValue::Decimal32(None, 5, 2)
501    )]
502    #[case::decimal64_some(
503        Scalar::decimal(
504            DecimalValue::I64(12345),
505            DecimalDType::new(10, 2),
506            Nullability::NonNullable
507        ),
508        ScalarValue::Decimal64(Some(12345), 10, 2)
509    )]
510    #[case::decimal64_null(
511        Scalar::null(DType::Decimal(DecimalDType::new(10, 2), Nullability::Nullable)),
512        ScalarValue::Decimal64(None, 10, 2)
513    )]
514    #[case::decimal128_some(
515        Scalar::decimal(
516            DecimalValue::I128(12345),
517            DecimalDType::new(20, 2),
518            Nullability::NonNullable
519        ),
520        ScalarValue::Decimal128(Some(12345), 20, 2)
521    )]
522    #[case::decimal128_null(
523        Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullability::Nullable)),
524        ScalarValue::Decimal128(None, 20, 2)
525    )]
526    #[case::decimal256_some(
527        Scalar::decimal(
528            DecimalValue::I256(i256::from(arrow_i256::from_i128(12345))),
529            DecimalDType::new(50, 10),
530            Nullability::NonNullable
531        ),
532        ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10)
533    )]
534    #[case::decimal256_null(
535        Scalar::null(DType::Decimal(DecimalDType::new(50, 10), Nullability::Nullable)),
536        ScalarValue::Decimal256(None, 50, 10)
537    )]
538    fn test_decimal_to_datafusion(
539        #[case] vortex_scalar: Scalar,
540        #[case] expected_df_scalar: ScalarValue,
541    ) {
542        let result = vortex_scalar.try_to_df().unwrap();
543        assert_eq!(result, expected_df_scalar);
544    }
545
546    #[rstest]
547    #[case::from_df_null(ScalarValue::Null, Scalar::null(DType::Null))]
548    #[case::from_df_bool_some(ScalarValue::Boolean(Some(true)), Scalar::from(true))]
549    #[case::from_df_bool_null(
550        ScalarValue::Boolean(None),
551        Scalar::null(DType::Bool(Nullability::Nullable))
552    )]
553    #[case::from_df_i32_some(ScalarValue::Int32(Some(42)), Scalar::from(42i32))]
554    #[case::from_df_i32_null(
555        ScalarValue::Int32(None),
556        Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))
557    )]
558    #[case::from_df_f64_some(ScalarValue::Float64(Some(2.5)), Scalar::from(2.5f64))]
559    #[case::from_df_f64_null(
560        ScalarValue::Float64(None),
561        Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))
562    )]
563    #[case::from_df_utf8_some(ScalarValue::Utf8(Some("test".to_string())), Scalar::from("test"))]
564    #[case::from_df_utf8_null(
565        ScalarValue::Utf8(None),
566        Scalar::null(DType::Utf8(Nullability::Nullable))
567    )]
568    #[case::from_df_binary_some(ScalarValue::Binary(Some(vec![1, 2, 3])), Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3]), Nullability::Nullable))]
569    #[case::from_df_binary_null(
570        ScalarValue::Binary(None),
571        Scalar::null(DType::Binary(Nullability::Nullable))
572    )]
573    fn test_from_datafusion_scalars(
574        #[case] df_scalar: ScalarValue,
575        #[case] expected_vortex: Scalar,
576    ) {
577        let result = Scalar::from_df(&df_scalar);
578        assert_eq!(result.dtype(), expected_vortex.dtype());
579        assert_eq!(result.is_null(), expected_vortex.is_null());
580
581        // For non-null values, convert both back to DataFusion for comparison
582        if !result.is_null() {
583            let result_df = result.try_to_df().unwrap();
584            let expected_df = expected_vortex.try_to_df().unwrap();
585            assert_eq!(result_df, expected_df);
586        }
587    }
588
589    #[rstest]
590    #[case::decimal128_some(ScalarValue::Decimal128(Some(12345), 10, 2))]
591    #[case::decimal128_null(ScalarValue::Decimal128(None, 10, 2))]
592    #[case::decimal256_some(ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10))]
593    #[case::decimal256_null(ScalarValue::Decimal256(None, 50, 10))]
594    fn test_from_datafusion_decimals(#[case] df_scalar: ScalarValue) {
595        let result = Scalar::from_df(&df_scalar);
596        match &df_scalar {
597            ScalarValue::Decimal128(value, precision, scale) => {
598                if let DType::Decimal(decimal_type, _) = result.dtype() {
599                    assert_eq!(decimal_type.precision(), *precision);
600                    assert_eq!(decimal_type.scale(), *scale);
601                    if value.is_some() {
602                        assert!(!result.is_null());
603                    } else {
604                        assert!(result.is_null());
605                    }
606                } else {
607                    panic!("Expected decimal type");
608                }
609            }
610            ScalarValue::Decimal256(value, precision, scale) => {
611                if let DType::Decimal(decimal_type, _) = result.dtype() {
612                    assert_eq!(decimal_type.precision(), *precision);
613                    assert_eq!(decimal_type.scale(), *scale);
614                    if value.is_some() {
615                        assert!(!result.is_null());
616                    } else {
617                        assert!(result.is_null());
618                    }
619                } else {
620                    panic!("Expected decimal type");
621                }
622            }
623            _ => panic!("Unexpected scalar type"),
624        }
625    }
626
627    #[rstest]
628    #[case::date32(ScalarValue::Date32(Some(18628)))] // 2021-01-01
629    #[case::date64(ScalarValue::Date64(Some(1609459200000)))] // 2021-01-01 in milliseconds
630    #[case::time32_second(ScalarValue::Time32Second(Some(3661)))] // 01:01:01
631    #[case::time32_millisecond(ScalarValue::Time32Millisecond(Some(3661000)))] // 01:01:01
632    #[case::time64_microsecond(ScalarValue::Time64Microsecond(Some(3661000000)))] // 01:01:01
633    #[case::time64_nanosecond(ScalarValue::Time64Nanosecond(Some(3661000000000)))] // 01:01:01
634    #[case::timestamp_second(ScalarValue::TimestampSecond(Some(1609459200), None))]
635    #[case::timestamp_millisecond(ScalarValue::TimestampMillisecond(Some(1609459200000), None))]
636    #[case::timestamp_microsecond(ScalarValue::TimestampMicrosecond(Some(1609459200000000), None))]
637    #[case::timestamp_nanosecond(ScalarValue::TimestampNanosecond(
638        Some(1609459200000000000),
639        None
640    ))]
641    fn test_from_datafusion_temporals(#[case] df_scalar: ScalarValue) {
642        let result = Scalar::from_df(&df_scalar);
643
644        // All temporal types should convert to extension types
645        if let DType::Extension(_) = result.dtype() {
646            assert!(!result.is_null());
647        } else {
648            panic!(
649                "Expected extension type for temporal scalar, got: {:?}",
650                result.dtype()
651            );
652        }
653    }
654
655    #[rstest]
656    #[case::u32(Scalar::from(42u32))]
657    #[case::i64(Scalar::from(-123i64))]
658    #[case::f64(Scalar::from(2.5f64))]
659    #[case::bool(Scalar::from(true))]
660    #[case::utf8(Scalar::from("hello world"))]
661    #[case::null_type(Scalar::null(DType::Null))]
662    #[case::null_i32(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)))]
663    #[case::decimal128(Scalar::decimal(
664        DecimalValue::I128(12345),
665        DecimalDType::new(10, 2),
666        Nullability::NonNullable
667    ))]
668    #[case::binary(Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3, 4, 5]), Nullability::NonNullable))]
669    fn test_round_trip_conversions(#[case] original: Scalar) {
670        let df_scalar = original.try_to_df().unwrap();
671        let round_trip = Scalar::from_df(&df_scalar);
672
673        // Check that core types match (ignoring nullability differences that can occur in round-trip)
674        assert!(
675            original.dtype().eq_ignore_nullability(round_trip.dtype()),
676            "DType mismatch for scalar: {:?} vs {:?}",
677            original.dtype(),
678            round_trip.dtype()
679        );
680
681        assert_eq!(
682            original.is_null(),
683            round_trip.is_null(),
684            "Null status mismatch for scalar: {:?}",
685            original
686        );
687
688        if !original.is_null() {
689            // For non-null values, compare by converting both to DataFusion scalars
690            let original_df = original.try_to_df().unwrap();
691            let round_trip_df = round_trip.try_to_df().unwrap();
692            assert_eq!(
693                original_df, round_trip_df,
694                "Value mismatch for scalar: {:?}",
695                original
696            );
697        }
698    }
699
700    #[rstest]
701    #[case::null_type(Scalar::null(DType::Null), ScalarValue::Null)]
702    #[case::null_bool(
703        Scalar::null(DType::Bool(Nullability::Nullable)),
704        ScalarValue::Boolean(None)
705    )]
706    #[case::null_i32(
707        Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
708        ScalarValue::Int32(None)
709    )]
710    #[case::null_f64(
711        Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)),
712        ScalarValue::Float64(None)
713    )]
714    #[case::null_utf8(
715        Scalar::null(DType::Utf8(Nullability::Nullable)),
716        ScalarValue::Utf8(None)
717    )]
718    #[case::null_binary(
719        Scalar::null(DType::Binary(Nullability::Nullable)),
720        ScalarValue::Binary(None)
721    )]
722    #[case::null_decimal128(
723        Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullability::Nullable)),
724        ScalarValue::Decimal128(None, 20, 2)
725    )]
726    #[case::null_decimal64(
727        Scalar::null(DType::Decimal(DecimalDType::new(10, 2), Nullability::Nullable)),
728        ScalarValue::Decimal64(None, 10, 2)
729    )]
730    #[case::null_decimal32(
731        Scalar::null(DType::Decimal(DecimalDType::new(5, 2), Nullability::Nullable)),
732        ScalarValue::Decimal32(None, 5, 2)
733    )]
734    fn test_null_handling(#[case] vortex_null: Scalar, #[case] expected_df_null: ScalarValue) {
735        // Test Vortex -> DataFusion
736        let df_result = vortex_null.try_to_df().unwrap();
737        assert_eq!(df_result, expected_df_null);
738
739        // Test DataFusion -> Vortex
740        let vortex_result = Scalar::from_df(&expected_df_null);
741        assert!(vortex_result.is_null());
742        assert!(
743            vortex_result
744                .dtype()
745                .eq_ignore_nullability(vortex_null.dtype())
746        );
747    }
748
749    #[rstest]
750    #[case::utf8(ScalarValue::Utf8(Some("test string".to_string())))]
751    #[case::utf8_view(ScalarValue::Utf8View(Some("test string".to_string())))]
752    #[case::large_utf8(ScalarValue::LargeUtf8(Some("test string".to_string())))]
753    fn test_utf8_variants(#[case] variant: ScalarValue) {
754        let result = Scalar::from_df(&variant);
755        assert_eq!(result.as_utf8().value().unwrap().as_str(), "test string");
756    }
757
758    #[rstest]
759    #[case::binary(ScalarValue::Binary(Some(vec![1u8, 2, 3, 4, 5])))]
760    #[case::binary_view(ScalarValue::BinaryView(Some(vec![1u8, 2, 3, 4, 5])))]
761    #[case::large_binary(ScalarValue::LargeBinary(Some(vec![1u8, 2, 3, 4, 5])))]
762    #[case::fixed_size_binary(ScalarValue::FixedSizeBinary(5, Some(vec![1u8, 2, 3, 4, 5])))]
763    fn test_binary_variants(#[case] variant: ScalarValue) {
764        let result = Scalar::from_df(&variant);
765        let result_bytes: Vec<u8> = result
766            .as_binary()
767            .value()
768            .cloned()
769            .unwrap()
770            .into_inner()
771            .into();
772        assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]);
773    }
774
775    #[test]
776    fn struct_scalar_round_trips() -> VortexResult<()> {
777        let dtype = DType::Struct(
778            StructFields::new(
779                FieldNames::from(["x", "y"]),
780                vec![
781                    DType::Primitive(PType::F64, Nullability::NonNullable),
782                    DType::Primitive(PType::F64, Nullability::NonNullable),
783                ],
784            ),
785            Nullability::NonNullable,
786        );
787        let original = Scalar::struct_(
788            dtype,
789            vec![Scalar::from(-111.7610f64), Scalar::from(34.8697f64)],
790        );
791
792        let df = original.try_to_df()?;
793        assert!(matches!(df, ScalarValue::Struct(_)));
794
795        // Back through `from_df` and out again yields the identical DataFusion struct value.
796        let back = Scalar::from_df(&df);
797        assert_eq!(back.try_to_df()?, df);
798        Ok(())
799    }
800
801    #[test]
802    fn null_struct_scalar_round_trips() -> VortexResult<()> {
803        let dtype = DType::Struct(
804            StructFields::new(
805                FieldNames::from(["x", "y"]),
806                vec![
807                    DType::Primitive(PType::F64, Nullability::Nullable),
808                    DType::Primitive(PType::F64, Nullability::Nullable),
809                ],
810            ),
811            Nullability::Nullable,
812        );
813
814        let df = Scalar::null(dtype).try_to_df()?;
815        assert!(matches!(df, ScalarValue::Struct(_)));
816        assert!(Scalar::from_df(&df).is_null());
817        Ok(())
818    }
819
820    #[rstest]
821    #[case::list(Scalar::null(DType::List(
822        Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
823        Nullability::Nullable
824    )))]
825    #[case::fixed_size_list(Scalar::null(DType::FixedSizeList(
826        Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
827        2,
828        Nullability::Nullable
829    )))]
830    #[case::union(Scalar::null(DType::Union(
831        UnionVariants::new(
832            ["a"].into(),
833            vec![DType::Primitive(PType::I32, Nullability::Nullable)],
834        )
835        .unwrap(),
836        Nullability::Nullable,
837    )))]
838    fn unsupported_vortex_scalars_return_errors(#[case] scalar: Scalar) {
839        let err = scalar.try_to_df().unwrap_err();
840
841        assert!(err.to_string().contains("unsupported scalar type"), "{err}");
842    }
843}
vortex_datafusion/convert/scalars.rs

vortex_datafusion/convert/
scalars.rs
