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::arrow::FromArrowType;
18use vortex::dtype::half::f16;
19use vortex::dtype::i256;
20use vortex::error::VortexExpect;
21use vortex::error::VortexResult;
22use vortex::error::vortex_bail;
23use vortex::error::vortex_err;
24use vortex::error::vortex_panic;
25use vortex::extension::datetime::AnyTemporal;
26use vortex::extension::datetime::TemporalMetadata;
27use vortex::extension::datetime::TimeUnit;
28use vortex::scalar::DecimalValue;
29use vortex::scalar::Scalar;
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::i256;
386    use vortex::scalar::DecimalValue;
387    use vortex::scalar::Scalar;
388
389    use super::*;
390
391    #[rstest]
392    #[case::u8_some(Scalar::from(42u8), ScalarValue::UInt8(Some(42)))]
393    #[case::u8_null(
394        Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)),
395        ScalarValue::UInt8(None)
396    )]
397    #[case::u16_some(Scalar::from(1234u16), ScalarValue::UInt16(Some(1234)))]
398    #[case::u16_null(
399        Scalar::null(DType::Primitive(PType::U16, Nullability::Nullable)),
400        ScalarValue::UInt16(None)
401    )]
402    #[case::u32_some(Scalar::from(123456u32), ScalarValue::UInt32(Some(123456)))]
403    #[case::u32_null(
404        Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)),
405        ScalarValue::UInt32(None)
406    )]
407    #[case::u64_some(Scalar::from(12345678u64), ScalarValue::UInt64(Some(12345678)))]
408    #[case::u64_null(
409        Scalar::null(DType::Primitive(PType::U64, Nullability::Nullable)),
410        ScalarValue::UInt64(None)
411    )]
412    #[case::i8_some(Scalar::from(-42i8), ScalarValue::Int8(Some(-42)))]
413    #[case::i8_null(
414        Scalar::null(DType::Primitive(PType::I8, Nullability::Nullable)),
415        ScalarValue::Int8(None)
416    )]
417    #[case::i16_some(Scalar::from(-1234i16), ScalarValue::Int16(Some(-1234)))]
418    #[case::i16_null(
419        Scalar::null(DType::Primitive(PType::I16, Nullability::Nullable)),
420        ScalarValue::Int16(None)
421    )]
422    #[case::i32_some(Scalar::from(-123456i32), ScalarValue::Int32(Some(-123456)))]
423    #[case::i32_null(
424        Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
425        ScalarValue::Int32(None)
426    )]
427    #[case::i64_some(Scalar::from(-12345678i64), ScalarValue::Int64(Some(-12345678)))]
428    #[case::i64_null(
429        Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable)),
430        ScalarValue::Int64(None)
431    )]
432    #[case::f32_some(Scalar::from(1.5f32), ScalarValue::Float32(Some(1.5)))]
433    #[case::f32_null(
434        Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable)),
435        ScalarValue::Float32(None)
436    )]
437    #[case::f64_some(Scalar::from(2.5f64), ScalarValue::Float64(Some(2.5)))]
438    #[case::f64_null(
439        Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)),
440        ScalarValue::Float64(None)
441    )]
442    fn test_primitive_to_datafusion(
443        #[case] vortex_scalar: Scalar,
444        #[case] expected_df_scalar: ScalarValue,
445    ) {
446        let result = vortex_scalar.try_to_df().unwrap();
447        assert_eq!(result, expected_df_scalar);
448    }
449
450    #[rstest]
451    #[case::bool_some(Scalar::from(true), ScalarValue::Boolean(Some(true)))]
452    #[case::bool_some_false(Scalar::from(false), ScalarValue::Boolean(Some(false)))]
453    #[case::bool_null(
454        Scalar::null(DType::Bool(Nullability::Nullable)),
455        ScalarValue::Boolean(None)
456    )]
457    #[case::null_type(Scalar::null(DType::Null), ScalarValue::Null)]
458    fn test_bool_and_null_to_datafusion(
459        #[case] vortex_scalar: Scalar,
460        #[case] expected_df_scalar: ScalarValue,
461    ) {
462        let result = vortex_scalar.try_to_df().unwrap();
463        assert_eq!(result, expected_df_scalar);
464    }
465
466    #[rstest]
467    #[case::utf8_some(Scalar::from("hello"), ScalarValue::Utf8(Some("hello".to_string())))]
468    #[case::utf8_null(
469        Scalar::null(DType::Utf8(Nullability::Nullable)),
470        ScalarValue::Utf8(None)
471    )]
472    #[case::binary_some(
473        Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3, 4]), Nullability::NonNullable),
474        ScalarValue::Binary(Some(vec![1u8, 2, 3, 4]))
475    )]
476    #[case::binary_null(
477        Scalar::null(DType::Binary(Nullability::Nullable)),
478        ScalarValue::Binary(None)
479    )]
480    fn test_string_and_binary_to_datafusion(
481        #[case] vortex_scalar: Scalar,
482        #[case] expected_df_scalar: ScalarValue,
483    ) {
484        let result = vortex_scalar.try_to_df().unwrap();
485        assert_eq!(result, expected_df_scalar);
486    }
487
488    #[rstest]
489    #[case::decimal32_some(
490        Scalar::decimal(
491            DecimalValue::I32(1234),
492            DecimalDType::new(5, 2),
493            Nullability::NonNullable
494        ),
495        ScalarValue::Decimal32(Some(1234), 5, 2)
496    )]
497    #[case::decimal32_null(
498        Scalar::null(DType::Decimal(DecimalDType::new(5, 2), Nullability::Nullable)),
499        ScalarValue::Decimal32(None, 5, 2)
500    )]
501    #[case::decimal64_some(
502        Scalar::decimal(
503            DecimalValue::I64(12345),
504            DecimalDType::new(10, 2),
505            Nullability::NonNullable
506        ),
507        ScalarValue::Decimal64(Some(12345), 10, 2)
508    )]
509    #[case::decimal64_null(
510        Scalar::null(DType::Decimal(DecimalDType::new(10, 2), Nullability::Nullable)),
511        ScalarValue::Decimal64(None, 10, 2)
512    )]
513    #[case::decimal128_some(
514        Scalar::decimal(
515            DecimalValue::I128(12345),
516            DecimalDType::new(20, 2),
517            Nullability::NonNullable
518        ),
519        ScalarValue::Decimal128(Some(12345), 20, 2)
520    )]
521    #[case::decimal128_null(
522        Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullability::Nullable)),
523        ScalarValue::Decimal128(None, 20, 2)
524    )]
525    #[case::decimal256_some(
526        Scalar::decimal(
527            DecimalValue::I256(i256::from(arrow_i256::from_i128(12345))),
528            DecimalDType::new(50, 10),
529            Nullability::NonNullable
530        ),
531        ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10)
532    )]
533    #[case::decimal256_null(
534        Scalar::null(DType::Decimal(DecimalDType::new(50, 10), Nullability::Nullable)),
535        ScalarValue::Decimal256(None, 50, 10)
536    )]
537    fn test_decimal_to_datafusion(
538        #[case] vortex_scalar: Scalar,
539        #[case] expected_df_scalar: ScalarValue,
540    ) {
541        let result = vortex_scalar.try_to_df().unwrap();
542        assert_eq!(result, expected_df_scalar);
543    }
544
545    #[rstest]
546    #[case::from_df_null(ScalarValue::Null, Scalar::null(DType::Null))]
547    #[case::from_df_bool_some(ScalarValue::Boolean(Some(true)), Scalar::from(true))]
548    #[case::from_df_bool_null(
549        ScalarValue::Boolean(None),
550        Scalar::null(DType::Bool(Nullability::Nullable))
551    )]
552    #[case::from_df_i32_some(ScalarValue::Int32(Some(42)), Scalar::from(42i32))]
553    #[case::from_df_i32_null(
554        ScalarValue::Int32(None),
555        Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))
556    )]
557    #[case::from_df_f64_some(ScalarValue::Float64(Some(2.5)), Scalar::from(2.5f64))]
558    #[case::from_df_f64_null(
559        ScalarValue::Float64(None),
560        Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))
561    )]
562    #[case::from_df_utf8_some(ScalarValue::Utf8(Some("test".to_string())), Scalar::from("test"))]
563    #[case::from_df_utf8_null(
564        ScalarValue::Utf8(None),
565        Scalar::null(DType::Utf8(Nullability::Nullable))
566    )]
567    #[case::from_df_binary_some(ScalarValue::Binary(Some(vec![1, 2, 3])), Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3]), Nullability::Nullable))]
568    #[case::from_df_binary_null(
569        ScalarValue::Binary(None),
570        Scalar::null(DType::Binary(Nullability::Nullable))
571    )]
572    fn test_from_datafusion_scalars(
573        #[case] df_scalar: ScalarValue,
574        #[case] expected_vortex: Scalar,
575    ) {
576        let result = Scalar::from_df(&df_scalar);
577        assert_eq!(result.dtype(), expected_vortex.dtype());
578        assert_eq!(result.is_null(), expected_vortex.is_null());
579
580        // For non-null values, convert both back to DataFusion for comparison
581        if !result.is_null() {
582            let result_df = result.try_to_df().unwrap();
583            let expected_df = expected_vortex.try_to_df().unwrap();
584            assert_eq!(result_df, expected_df);
585        }
586    }
587
588    #[rstest]
589    #[case::decimal128_some(ScalarValue::Decimal128(Some(12345), 10, 2))]
590    #[case::decimal128_null(ScalarValue::Decimal128(None, 10, 2))]
591    #[case::decimal256_some(ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10))]
592    #[case::decimal256_null(ScalarValue::Decimal256(None, 50, 10))]
593    fn test_from_datafusion_decimals(#[case] df_scalar: ScalarValue) {
594        let result = Scalar::from_df(&df_scalar);
595        match &df_scalar {
596            ScalarValue::Decimal128(value, precision, scale) => {
597                if let DType::Decimal(decimal_type, _) = result.dtype() {
598                    assert_eq!(decimal_type.precision(), *precision);
599                    assert_eq!(decimal_type.scale(), *scale);
600                    if value.is_some() {
601                        assert!(!result.is_null());
602                    } else {
603                        assert!(result.is_null());
604                    }
605                } else {
606                    panic!("Expected decimal type");
607                }
608            }
609            ScalarValue::Decimal256(value, precision, scale) => {
610                if let DType::Decimal(decimal_type, _) = result.dtype() {
611                    assert_eq!(decimal_type.precision(), *precision);
612                    assert_eq!(decimal_type.scale(), *scale);
613                    if value.is_some() {
614                        assert!(!result.is_null());
615                    } else {
616                        assert!(result.is_null());
617                    }
618                } else {
619                    panic!("Expected decimal type");
620                }
621            }
622            _ => panic!("Unexpected scalar type"),
623        }
624    }
625
626    #[rstest]
627    #[case::date32(ScalarValue::Date32(Some(18628)))] // 2021-01-01
628    #[case::date64(ScalarValue::Date64(Some(1609459200000)))] // 2021-01-01 in milliseconds
629    #[case::time32_second(ScalarValue::Time32Second(Some(3661)))] // 01:01:01
630    #[case::time32_millisecond(ScalarValue::Time32Millisecond(Some(3661000)))] // 01:01:01
631    #[case::time64_microsecond(ScalarValue::Time64Microsecond(Some(3661000000)))] // 01:01:01
632    #[case::time64_nanosecond(ScalarValue::Time64Nanosecond(Some(3661000000000)))] // 01:01:01
633    #[case::timestamp_second(ScalarValue::TimestampSecond(Some(1609459200), None))]
634    #[case::timestamp_millisecond(ScalarValue::TimestampMillisecond(Some(1609459200000), None))]
635    #[case::timestamp_microsecond(ScalarValue::TimestampMicrosecond(Some(1609459200000000), None))]
636    #[case::timestamp_nanosecond(ScalarValue::TimestampNanosecond(
637        Some(1609459200000000000),
638        None
639    ))]
640    fn test_from_datafusion_temporals(#[case] df_scalar: ScalarValue) {
641        let result = Scalar::from_df(&df_scalar);
642
643        // All temporal types should convert to extension types
644        if let DType::Extension(_) = result.dtype() {
645            assert!(!result.is_null());
646        } else {
647            panic!(
648                "Expected extension type for temporal scalar, got: {:?}",
649                result.dtype()
650            );
651        }
652    }
653
654    #[rstest]
655    #[case::u32(Scalar::from(42u32))]
656    #[case::i64(Scalar::from(-123i64))]
657    #[case::f64(Scalar::from(2.5f64))]
658    #[case::bool(Scalar::from(true))]
659    #[case::utf8(Scalar::from("hello world"))]
660    #[case::null_type(Scalar::null(DType::Null))]
661    #[case::null_i32(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)))]
662    #[case::decimal128(Scalar::decimal(
663        DecimalValue::I128(12345),
664        DecimalDType::new(10, 2),
665        Nullability::NonNullable
666    ))]
667    #[case::binary(Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3, 4, 5]), Nullability::NonNullable))]
668    fn test_round_trip_conversions(#[case] original: Scalar) {
669        let df_scalar = original.try_to_df().unwrap();
670        let round_trip = Scalar::from_df(&df_scalar);
671
672        // Check that core types match (ignoring nullability differences that can occur in round-trip)
673        assert!(
674            original.dtype().eq_ignore_nullability(round_trip.dtype()),
675            "DType mismatch for scalar: {:?} vs {:?}",
676            original.dtype(),
677            round_trip.dtype()
678        );
679
680        assert_eq!(
681            original.is_null(),
682            round_trip.is_null(),
683            "Null status mismatch for scalar: {:?}",
684            original
685        );
686
687        if !original.is_null() {
688            // For non-null values, compare by converting both to DataFusion scalars
689            let original_df = original.try_to_df().unwrap();
690            let round_trip_df = round_trip.try_to_df().unwrap();
691            assert_eq!(
692                original_df, round_trip_df,
693                "Value mismatch for scalar: {:?}",
694                original
695            );
696        }
697    }
698
699    #[rstest]
700    #[case::null_type(Scalar::null(DType::Null), ScalarValue::Null)]
701    #[case::null_bool(
702        Scalar::null(DType::Bool(Nullability::Nullable)),
703        ScalarValue::Boolean(None)
704    )]
705    #[case::null_i32(
706        Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
707        ScalarValue::Int32(None)
708    )]
709    #[case::null_f64(
710        Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)),
711        ScalarValue::Float64(None)
712    )]
713    #[case::null_utf8(
714        Scalar::null(DType::Utf8(Nullability::Nullable)),
715        ScalarValue::Utf8(None)
716    )]
717    #[case::null_binary(
718        Scalar::null(DType::Binary(Nullability::Nullable)),
719        ScalarValue::Binary(None)
720    )]
721    #[case::null_decimal128(
722        Scalar::null(DType::Decimal(DecimalDType::new(20, 2), Nullability::Nullable)),
723        ScalarValue::Decimal128(None, 20, 2)
724    )]
725    #[case::null_decimal64(
726        Scalar::null(DType::Decimal(DecimalDType::new(10, 2), Nullability::Nullable)),
727        ScalarValue::Decimal64(None, 10, 2)
728    )]
729    #[case::null_decimal32(
730        Scalar::null(DType::Decimal(DecimalDType::new(5, 2), Nullability::Nullable)),
731        ScalarValue::Decimal32(None, 5, 2)
732    )]
733    fn test_null_handling(#[case] vortex_null: Scalar, #[case] expected_df_null: ScalarValue) {
734        // Test Vortex -> DataFusion
735        let df_result = vortex_null.try_to_df().unwrap();
736        assert_eq!(df_result, expected_df_null);
737
738        // Test DataFusion -> Vortex
739        let vortex_result = Scalar::from_df(&expected_df_null);
740        assert!(vortex_result.is_null());
741        assert!(
742            vortex_result
743                .dtype()
744                .eq_ignore_nullability(vortex_null.dtype())
745        );
746    }
747
748    #[rstest]
749    #[case::utf8(ScalarValue::Utf8(Some("test string".to_string())))]
750    #[case::utf8_view(ScalarValue::Utf8View(Some("test string".to_string())))]
751    #[case::large_utf8(ScalarValue::LargeUtf8(Some("test string".to_string())))]
752    fn test_utf8_variants(#[case] variant: ScalarValue) {
753        let result = Scalar::from_df(&variant);
754        assert_eq!(result.as_utf8().value().unwrap().as_str(), "test string");
755    }
756
757    #[rstest]
758    #[case::binary(ScalarValue::Binary(Some(vec![1u8, 2, 3, 4, 5])))]
759    #[case::binary_view(ScalarValue::BinaryView(Some(vec![1u8, 2, 3, 4, 5])))]
760    #[case::large_binary(ScalarValue::LargeBinary(Some(vec![1u8, 2, 3, 4, 5])))]
761    #[case::fixed_size_binary(ScalarValue::FixedSizeBinary(5, Some(vec![1u8, 2, 3, 4, 5])))]
762    fn test_binary_variants(#[case] variant: ScalarValue) {
763        let result = Scalar::from_df(&variant);
764        let result_bytes: Vec<u8> = result
765            .as_binary()
766            .value()
767            .cloned()
768            .unwrap()
769            .into_inner()
770            .into();
771        assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]);
772    }
773
774    #[test]
775    fn struct_scalar_round_trips() -> VortexResult<()> {
776        let dtype = DType::Struct(
777            StructFields::new(
778                FieldNames::from(["x", "y"]),
779                vec![
780                    DType::Primitive(PType::F64, Nullability::NonNullable),
781                    DType::Primitive(PType::F64, Nullability::NonNullable),
782                ],
783            ),
784            Nullability::NonNullable,
785        );
786        let original = Scalar::struct_(
787            dtype,
788            vec![Scalar::from(-111.7610f64), Scalar::from(34.8697f64)],
789        );
790
791        let df = original.try_to_df()?;
792        assert!(matches!(df, ScalarValue::Struct(_)));
793
794        // Back through `from_df` and out again yields the identical DataFusion struct value.
795        let back = Scalar::from_df(&df);
796        assert_eq!(back.try_to_df()?, df);
797        Ok(())
798    }
799
800    #[test]
801    fn null_struct_scalar_round_trips() -> VortexResult<()> {
802        let dtype = DType::Struct(
803            StructFields::new(
804                FieldNames::from(["x", "y"]),
805                vec![
806                    DType::Primitive(PType::F64, Nullability::Nullable),
807                    DType::Primitive(PType::F64, Nullability::Nullable),
808                ],
809            ),
810            Nullability::Nullable,
811        );
812
813        let df = Scalar::null(dtype).try_to_df()?;
814        assert!(matches!(df, ScalarValue::Struct(_)));
815        assert!(Scalar::from_df(&df).is_null());
816        Ok(())
817    }
818
819    #[rstest]
820    #[case::list(Scalar::null(DType::List(
821        Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
822        Nullability::Nullable
823    )))]
824    #[case::fixed_size_list(Scalar::null(DType::FixedSizeList(
825        Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
826        2,
827        Nullability::Nullable
828    )))]
829    #[case::union(Scalar::null(DType::Union(Nullability::Nullable)))]
830    fn unsupported_vortex_scalars_return_errors(#[case] scalar: Scalar) {
831        let err = scalar.try_to_df().unwrap_err();
832
833        assert!(err.to_string().contains("unsupported scalar type"), "{err}");
834    }
835}
vortex_datafusion/convert/scalars.rs

vortex_datafusion/convert/
scalars.rs
