vortex_arrow/
scalar.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Conversions between [`Scalar`] and Arrow scalar types.
5
6use std::sync::Arc;
7
8use arrow_array::Scalar as ArrowScalar;
9use arrow_array::*;
10use arrow_buffer::NullBuffer;
11use arrow_buffer::OffsetBuffer;
12use arrow_schema::Field;
13use arrow_schema::Fields;
14use vortex_array::dtype::DType;
15use vortex_array::dtype::PType;
16use vortex_array::dtype::i256;
17use vortex_array::extension::datetime::AnyTemporal;
18use vortex_array::extension::datetime::TemporalMetadata;
19use vortex_array::extension::datetime::TimeUnit;
20use vortex_array::scalar::BinaryScalar;
21use vortex_array::scalar::BoolScalar;
22use vortex_array::scalar::DecimalScalar;
23use vortex_array::scalar::ExtScalar;
24use vortex_array::scalar::MapScalar;
25use vortex_array::scalar::PrimitiveScalar;
26use vortex_array::scalar::Scalar;
27use vortex_array::scalar::Utf8Scalar;
28use vortex_error::VortexError;
29use vortex_error::vortex_bail;
30use vortex_error::vortex_err;
31
32use crate::dtype::to_data_type_naive;
33
34/// Arrow represents scalars as single-element arrays. This constant is the length of those arrays.
35const SCALAR_ARRAY_LEN: usize = 1;
36
37/// Converts an optional value to an Arrow scalar array.
38macro_rules! value_to_arrow_scalar {
39    ($V:expr, $AR:ty) => {
40        Ok(std::sync::Arc::new(
41            $V.map(<$AR>::new_scalar)
42                .unwrap_or_else(|| arrow_array::Scalar::new(<$AR>::new_null(SCALAR_ARRAY_LEN))),
43        ))
44    };
45}
46
47/// Converts an optional timestamp value to an Arrow scalar array.
48macro_rules! timestamp_to_arrow_scalar {
49    ($V:expr, $TZ:expr, $AR:ty) => {{
50        let array = match $V {
51            Some(v) => <$AR>::new_scalar(v).into_inner(),
52            None => <$AR>::new_null(SCALAR_ARRAY_LEN),
53        }
54        .with_timezone_opt($TZ);
55        Ok(Arc::new(ArrowScalar::new(array)))
56    }};
57}
58
59/// Convert a Vortex [`Scalar`] into an Arrow [`Datum`] (a single-element Arrow array).
60///
61/// This mirrors a `TryFrom<&Scalar> for Arc<dyn Datum>` conversion; a separate trait is
62/// required because both `Scalar` and `Datum` are foreign to this crate.
63pub trait ToArrowDatum {
64    /// Convert this scalar to an Arrow [`Datum`].
65    fn to_arrow_datum(&self) -> Result<Arc<dyn Datum>, VortexError>;
66}
67
68impl ToArrowDatum for Scalar {
69    fn to_arrow_datum(&self) -> Result<Arc<dyn Datum>, VortexError> {
70        let value = self;
71        match value.dtype() {
72            DType::Null => Ok(Arc::new(NullArray::new(SCALAR_ARRAY_LEN))),
73            DType::Bool(_) => bool_to_arrow(value.as_bool()),
74            DType::Primitive(..) => primitive_to_arrow(value.as_primitive()),
75            DType::Decimal(..) => decimal_to_arrow(value.as_decimal()),
76            DType::Utf8(_) => utf8_to_arrow(value.as_utf8()),
77            DType::Binary(_) => binary_to_arrow(value.as_binary()),
78            DType::List(..) => vortex_bail!("list scalar conversion is not supported"),
79            DType::FixedSizeList(..) => {
80                vortex_bail!("fixed-size list scalar conversion is not supported")
81            }
82            DType::Map(..) => map_to_arrow(value.as_map()),
83            DType::Struct(..) => vortex_bail!("struct scalar conversion is not supported"),
84            DType::Union(..) => vortex_bail!("union scalar conversion is not supported"),
85            DType::Variant(_) => vortex_bail!("Variant scalar conversion is not supported"),
86            DType::Extension(..) => extension_to_arrow(value.as_extension()),
87        }
88    }
89}
90
91/// Convert a [`BoolScalar`] to an Arrow [`Datum`].
92fn bool_to_arrow(scalar: BoolScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
93    value_to_arrow_scalar!(scalar.value(), BooleanArray)
94}
95
96/// Convert a [`PrimitiveScalar`] to an Arrow [`Datum`].
97fn primitive_to_arrow(scalar: PrimitiveScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
98    match scalar.ptype() {
99        PType::U8 => value_to_arrow_scalar!(scalar.typed_value(), UInt8Array),
100        PType::U16 => value_to_arrow_scalar!(scalar.typed_value(), UInt16Array),
101        PType::U32 => value_to_arrow_scalar!(scalar.typed_value(), UInt32Array),
102        PType::U64 => value_to_arrow_scalar!(scalar.typed_value(), UInt64Array),
103        PType::I8 => value_to_arrow_scalar!(scalar.typed_value(), Int8Array),
104        PType::I16 => value_to_arrow_scalar!(scalar.typed_value(), Int16Array),
105        PType::I32 => value_to_arrow_scalar!(scalar.typed_value(), Int32Array),
106        PType::I64 => value_to_arrow_scalar!(scalar.typed_value(), Int64Array),
107        PType::F16 => value_to_arrow_scalar!(scalar.typed_value(), Float16Array),
108        PType::F32 => value_to_arrow_scalar!(scalar.typed_value(), Float32Array),
109        PType::F64 => value_to_arrow_scalar!(scalar.typed_value(), Float64Array),
110    }
111}
112
113/// Convert a [`DecimalScalar`] to an Arrow [`Datum`].
114fn decimal_to_arrow(scalar: DecimalScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
115    let DType::Decimal(decimal_dtype, _) = scalar.dtype() else {
116        vortex_bail!("Expected decimal scalar, got {}", scalar.dtype());
117    };
118    let precision = decimal_dtype.precision();
119    let scale = decimal_dtype.scale();
120    // TODO(joe): Replace with decimal32, etc. once Arrow supports them.
121    match scalar.decimal_value() {
122        Some(value) => {
123            let value = value.as_i256();
124            if precision <= 38 {
125                let value = value.maybe_i128().ok_or_else(|| {
126                    vortex_err!(
127                        "Decimal value {value} cannot fit in Arrow Decimal128 for precision {precision}"
128                    )
129                })?;
130                decimal128_scalar(value, precision, scale)
131            } else {
132                decimal256_scalar(value, precision, scale)
133            }
134        }
135        None => {
136            let data_type = to_data_type_naive(scalar.dtype())?;
137            Ok(Arc::new(ArrowScalar::new(new_null_array(
138                &data_type,
139                SCALAR_ARRAY_LEN,
140            ))))
141        }
142    }
143}
144
145fn decimal128_scalar(value: i128, precision: u8, scale: i8) -> Result<Arc<dyn Datum>, VortexError> {
146    let array = Decimal128Array::new_scalar(value)
147        .into_inner()
148        .with_precision_and_scale(precision, scale)?;
149    Ok(Arc::new(ArrowScalar::new(array)))
150}
151
152fn decimal256_scalar(value: i256, precision: u8, scale: i8) -> Result<Arc<dyn Datum>, VortexError> {
153    let array = Decimal256Array::new_scalar(value.into())
154        .into_inner()
155        .with_precision_and_scale(precision, scale)?;
156    Ok(Arc::new(ArrowScalar::new(array)))
157}
158
159/// Convert a [`Utf8Scalar`] to an Arrow [`Datum`].
160fn utf8_to_arrow(scalar: Utf8Scalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
161    value_to_arrow_scalar!(scalar.value(), StringViewArray)
162}
163
164/// Convert a [`BinaryScalar`] to an Arrow [`Datum`].
165fn binary_to_arrow(scalar: BinaryScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
166    value_to_arrow_scalar!(scalar.value(), BinaryViewArray)
167}
168
169/// Convert a [`MapScalar`] to an Arrow [`Datum`].
170fn map_to_arrow(scalar: MapScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
171    let map_dtype = scalar.map_dtype();
172    let key_dtype = map_dtype.key_dtype();
173    let value_dtype = map_dtype.value_dtype();
174    let key_field = Field::new("key", to_data_type_naive(&key_dtype)?, false);
175    let value_field = Field::new(
176        "value",
177        to_data_type_naive(&value_dtype)?,
178        value_dtype.is_nullable(),
179    );
180    let fields = Fields::from(vec![key_field, value_field]);
181
182    let entries = scalar.entries().collect::<Vec<_>>();
183    let keys = entries
184        .iter()
185        .map(|(key, _)| key.to_arrow_datum())
186        .collect::<Result<Vec<_>, _>>()?;
187    let values = entries
188        .iter()
189        .map(|(_, value)| value.to_arrow_datum())
190        .collect::<Result<Vec<_>, _>>()?;
191
192    let key_array = concat_scalar_arrays(&keys, &key_dtype)?;
193    let value_array = concat_scalar_arrays(&values, &value_dtype)?;
194    let entries = StructArray::try_new_with_length(
195        fields.clone(),
196        vec![key_array, value_array],
197        None,
198        entries.len(),
199    )?;
200
201    let entries_len = entries.len();
202    let entries_len = i32::try_from(entries_len).map_err(|_| {
203        vortex_err!(
204            "Cannot convert map scalar with {entries_len} entries to Arrow: MapArray offsets are i32"
205        )
206    })?;
207    let offsets = OffsetBuffer::new(vec![0_i32, entries_len].into());
208    let entries_field = Arc::new(Field::new_struct("entries", fields, false));
209    let nulls = scalar
210        .is_null()
211        .then(|| NullBuffer::new_null(SCALAR_ARRAY_LEN));
212    let map = MapArray::try_new(
213        entries_field,
214        offsets,
215        entries,
216        nulls,
217        map_dtype.keys_sorted(),
218    )?;
219    Ok(Arc::new(ArrowScalar::new(map)))
220}
221
222fn concat_scalar_arrays(
223    scalars: &[Arc<dyn Datum>],
224    dtype: &DType,
225) -> Result<ArrayRef, VortexError> {
226    if scalars.is_empty() {
227        return Ok(new_empty_array(&to_data_type_naive(dtype)?));
228    }
229
230    let arrays = scalars
231        .iter()
232        .map(|scalar| scalar.get().0)
233        .collect::<Vec<_>>();
234    Ok(arrow_select::concat::concat(&arrays)?)
235}
236
237/// Convert an [`ExtScalar`] to an Arrow [`Datum`].
238///
239/// Currently only temporal extension types (timestamps, dates, and times) are supported.
240fn extension_to_arrow(scalar: ExtScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
241    let ext_dtype = scalar.ext_dtype();
242    let Some(temporal) = ext_dtype.metadata_opt::<AnyTemporal>() else {
243        vortex_bail!(
244            "Cannot convert extension scalar {} to Arrow",
245            ext_dtype.id()
246        )
247    };
248
249    let storage_scalar = scalar.to_storage_scalar();
250    let primitive = storage_scalar
251        .as_primitive_opt()
252        .ok_or_else(|| vortex_err!("Expected primitive scalar"))?;
253
254    match temporal {
255        TemporalMetadata::Timestamp(unit, tz) => {
256            let value = primitive.as_::<i64>();
257            match unit {
258                TimeUnit::Nanoseconds => {
259                    timestamp_to_arrow_scalar!(value, tz.clone(), TimestampNanosecondArray)
260                }
261                TimeUnit::Microseconds => {
262                    timestamp_to_arrow_scalar!(value, tz.clone(), TimestampMicrosecondArray)
263                }
264                TimeUnit::Milliseconds => {
265                    timestamp_to_arrow_scalar!(value, tz.clone(), TimestampMillisecondArray)
266                }
267                TimeUnit::Seconds => {
268                    timestamp_to_arrow_scalar!(value, tz.clone(), TimestampSecondArray)
269                }
270                TimeUnit::Days => {
271                    vortex_bail!("Unsupported TimeUnit {unit} for {}", ext_dtype.id())
272                }
273            }
274        }
275        TemporalMetadata::Date(unit) => match unit {
276            TimeUnit::Milliseconds => {
277                value_to_arrow_scalar!(primitive.as_::<i64>(), Date64Array)
278            }
279            TimeUnit::Days => {
280                value_to_arrow_scalar!(primitive.as_::<i32>(), Date32Array)
281            }
282            TimeUnit::Nanoseconds | TimeUnit::Microseconds | TimeUnit::Seconds => {
283                vortex_bail!("Unsupported TimeUnit {unit} for {}", ext_dtype.id())
284            }
285        },
286        TemporalMetadata::Time(unit) => match unit {
287            TimeUnit::Nanoseconds => {
288                value_to_arrow_scalar!(primitive.as_::<i64>(), Time64NanosecondArray)
289            }
290            TimeUnit::Microseconds => {
291                value_to_arrow_scalar!(primitive.as_::<i64>(), Time64MicrosecondArray)
292            }
293            TimeUnit::Milliseconds => {
294                value_to_arrow_scalar!(primitive.as_::<i32>(), Time32MillisecondArray)
295            }
296            TimeUnit::Seconds => {
297                value_to_arrow_scalar!(primitive.as_::<i32>(), Time32SecondArray)
298            }
299            TimeUnit::Days => {
300                vortex_bail!("Unsupported TimeUnit {unit} for {}", ext_dtype.id())
301            }
302        },
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use std::sync::Arc;
309
310    use arrow_array::Array;
311    use arrow_array::Decimal128Array;
312    use arrow_array::Int32Array;
313    use arrow_array::MapArray;
314    use arrow_array::StringViewArray;
315    use arrow_schema::DataType;
316    use rstest::rstest;
317    use vortex_array::dtype::DType;
318    use vortex_array::dtype::DecimalDType;
319    use vortex_array::dtype::FieldDType;
320    use vortex_array::dtype::NativeDType;
321    use vortex_array::dtype::Nullability;
322    use vortex_array::dtype::PType;
323    use vortex_array::dtype::StructFields;
324    use vortex_array::dtype::extension::ExtDType;
325    use vortex_array::dtype::extension::ExtId;
326    use vortex_array::dtype::extension::ExtVTable;
327    use vortex_array::dtype::i256;
328    use vortex_array::extension::datetime::Date;
329    use vortex_array::extension::datetime::Time;
330    use vortex_array::extension::datetime::TimeUnit;
331    use vortex_array::extension::datetime::Timestamp;
332    use vortex_array::extension::datetime::TimestampOptions;
333    use vortex_array::scalar::DecimalValue;
334    use vortex_array::scalar::Scalar;
335    use vortex_array::scalar::ScalarValue;
336    use vortex_error::VortexResult;
337    use vortex_error::vortex_bail;
338
339    use super::ToArrowDatum;
340
341    #[test]
342    fn test_null_scalar_to_arrow() {
343        let scalar = Scalar::null(DType::Null);
344        let result = scalar.to_arrow_datum();
345        assert!(result.is_ok());
346    }
347
348    #[test]
349    fn test_bool_scalar_to_arrow() {
350        let scalar = Scalar::bool(true, Nullability::NonNullable);
351        let result = scalar.to_arrow_datum();
352        assert!(result.is_ok());
353    }
354
355    #[test]
356    fn test_null_bool_scalar_to_arrow() {
357        let scalar = Scalar::null(bool::dtype().as_nullable());
358        let result = scalar.to_arrow_datum();
359        assert!(result.is_ok());
360    }
361
362    #[test]
363    fn test_primitive_u8_to_arrow() {
364        let scalar = Scalar::primitive(42u8, Nullability::NonNullable);
365        let result = scalar.to_arrow_datum();
366        assert!(result.is_ok());
367    }
368
369    #[test]
370    fn test_primitive_u16_to_arrow() {
371        let scalar = Scalar::primitive(1000u16, Nullability::NonNullable);
372        let result = scalar.to_arrow_datum();
373        assert!(result.is_ok());
374    }
375
376    #[test]
377    fn test_primitive_u32_to_arrow() {
378        let scalar = Scalar::primitive(100000u32, Nullability::NonNullable);
379        let result = scalar.to_arrow_datum();
380        assert!(result.is_ok());
381    }
382
383    #[test]
384    fn test_primitive_u64_to_arrow() {
385        let scalar = Scalar::primitive(10000000000u64, Nullability::NonNullable);
386        let result = scalar.to_arrow_datum();
387        assert!(result.is_ok());
388    }
389
390    #[test]
391    fn test_primitive_i8_to_arrow() {
392        let scalar = Scalar::primitive(-42i8, Nullability::NonNullable);
393        let result = scalar.to_arrow_datum();
394        assert!(result.is_ok());
395    }
396
397    #[test]
398    fn test_primitive_i16_to_arrow() {
399        let scalar = Scalar::primitive(-1000i16, Nullability::NonNullable);
400        let result = scalar.to_arrow_datum();
401        assert!(result.is_ok());
402    }
403
404    #[test]
405    fn test_primitive_i32_to_arrow() {
406        let scalar = Scalar::primitive(-100000i32, Nullability::NonNullable);
407        let result = scalar.to_arrow_datum();
408        assert!(result.is_ok());
409    }
410
411    #[test]
412    fn test_primitive_i64_to_arrow() {
413        let scalar = Scalar::primitive(-10000000000i64, Nullability::NonNullable);
414        let result = scalar.to_arrow_datum();
415        assert!(result.is_ok());
416    }
417
418    #[test]
419    fn test_primitive_f16_to_arrow() {
420        use vortex_array::dtype::half::f16;
421
422        let scalar = Scalar::primitive(f16::from_f32(1.234), Nullability::NonNullable);
423        let result = scalar.to_arrow_datum();
424        assert!(result.is_ok());
425    }
426
427    #[test]
428    fn test_primitive_f32_to_arrow() {
429        let scalar = Scalar::primitive(1.234f32, Nullability::NonNullable);
430        let result = scalar.to_arrow_datum();
431        assert!(result.is_ok());
432    }
433
434    #[test]
435    fn test_primitive_f64_to_arrow() {
436        let scalar = Scalar::primitive(1.234567890123f64, Nullability::NonNullable);
437        let result = scalar.to_arrow_datum();
438        assert!(result.is_ok());
439    }
440
441    #[test]
442    fn test_null_primitive_to_arrow() {
443        let scalar = Scalar::null(i32::dtype().as_nullable());
444        let result = scalar.to_arrow_datum();
445        assert!(result.is_ok());
446    }
447
448    #[test]
449    fn test_utf8_scalar_to_arrow() {
450        let scalar = Scalar::utf8("hello world".to_string(), Nullability::NonNullable);
451        let result = scalar.to_arrow_datum();
452        assert!(result.is_ok());
453    }
454
455    #[test]
456    fn test_null_utf8_scalar_to_arrow() {
457        let scalar = Scalar::null(String::dtype().as_nullable());
458        let result = scalar.to_arrow_datum();
459        assert!(result.is_ok());
460    }
461
462    #[test]
463    fn test_binary_scalar_to_arrow() {
464        let data = vec![1u8, 2, 3, 4, 5];
465        let scalar = Scalar::binary(data, Nullability::NonNullable);
466        let result = scalar.to_arrow_datum();
467        assert!(result.is_ok());
468    }
469
470    #[test]
471    fn test_null_binary_scalar_to_arrow() {
472        let scalar = Scalar::null(DType::Binary(Nullability::Nullable));
473        let result = scalar.to_arrow_datum();
474        assert!(result.is_ok());
475    }
476
477    fn assert_arrow_scalar_data_type(scalar: &Scalar, expected: DataType) -> VortexResult<()> {
478        let datum = scalar.to_arrow_datum()?;
479        let (array, is_scalar) = datum.get();
480        assert!(is_scalar);
481        assert_eq!(array.data_type(), &expected);
482        Ok(())
483    }
484
485    #[test]
486    fn test_decimal_scalars_to_arrow() -> VortexResult<()> {
487        // Test various decimal value types
488        let decimal_dtype = DecimalDType::new(5, 2);
489
490        let scalar_i8 = Scalar::decimal(
491            DecimalValue::I8(100),
492            decimal_dtype,
493            Nullability::NonNullable,
494        );
495        assert_arrow_scalar_data_type(&scalar_i8, DataType::Decimal128(5, 2))?;
496
497        let scalar_i16 = Scalar::decimal(
498            DecimalValue::I16(10000),
499            decimal_dtype,
500            Nullability::NonNullable,
501        );
502        assert_arrow_scalar_data_type(&scalar_i16, DataType::Decimal128(5, 2))?;
503
504        let scalar_i32 = Scalar::decimal(
505            DecimalValue::I32(99999),
506            decimal_dtype,
507            Nullability::NonNullable,
508        );
509        assert_arrow_scalar_data_type(&scalar_i32, DataType::Decimal128(5, 2))?;
510
511        let scalar_i64 = Scalar::decimal(
512            DecimalValue::I64(99999),
513            decimal_dtype,
514            Nullability::NonNullable,
515        );
516        assert_arrow_scalar_data_type(&scalar_i64, DataType::Decimal128(5, 2))?;
517
518        let scalar_i128 = Scalar::decimal(
519            DecimalValue::I128(99999),
520            decimal_dtype,
521            Nullability::NonNullable,
522        );
523        assert_arrow_scalar_data_type(&scalar_i128, DataType::Decimal128(5, 2))?;
524
525        let value_i256 = i256::from_i128(99999);
526        let scalar_i256 = Scalar::decimal(
527            DecimalValue::I256(value_i256),
528            decimal_dtype,
529            Nullability::NonNullable,
530        );
531        assert_arrow_scalar_data_type(&scalar_i256, DataType::Decimal128(5, 2))?;
532
533        Ok(())
534    }
535
536    #[test]
537    fn decimal_i64_with_wide_precision_exports_decimal256() -> VortexResult<()> {
538        let scalar = Scalar::decimal(
539            DecimalValue::I64(1),
540            DecimalDType::new(39, 0),
541            Nullability::NonNullable,
542        );
543
544        assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 0))
545    }
546
547    #[test]
548    fn decimal_i256_with_narrow_precision_exports_decimal128() -> VortexResult<()> {
549        let scalar = Scalar::decimal(
550            DecimalValue::I256(i256::from_i128(1234)),
551            DecimalDType::new(4, 2),
552            Nullability::NonNullable,
553        );
554
555        assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(4, 2))
556    }
557
558    #[test]
559    fn test_null_decimal_to_arrow() -> VortexResult<()> {
560        let decimal_dtype = DecimalDType::new(10, 2);
561        let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable));
562        assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(10, 2))?;
563
564        let decimal_dtype = DecimalDType::new(39, 2);
565        let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable));
566        assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 2))
567    }
568
569    #[test]
570    fn decimal_scalar_to_arrow_preserves_precision_and_scale() -> VortexResult<()> {
571        let decimal_dtype = DecimalDType::new(12, 3);
572        let scalar = Scalar::decimal(
573            DecimalValue::I128(12345),
574            decimal_dtype,
575            Nullability::NonNullable,
576        );
577
578        let datum = scalar.to_arrow_datum()?;
579        let (array, is_scalar) = datum.get();
580        assert!(is_scalar);
581        let decimal = array
582            .as_any()
583            .downcast_ref::<Decimal128Array>()
584            .expect("decimal scalar should convert to Decimal128");
585        assert_eq!(decimal.precision(), 12);
586        assert_eq!(decimal.scale(), 3);
587
588        Ok(())
589    }
590
591    #[test]
592    fn test_map_scalar_to_arrow() -> VortexResult<()> {
593        let dtype = DType::map(
594            DType::Primitive(PType::I32, Nullability::NonNullable),
595            DType::Utf8(Nullability::Nullable),
596            true,
597            Nullability::Nullable,
598        )?;
599        let scalar = Scalar::try_map(
600            dtype,
601            [
602                (
603                    Scalar::primitive(1i32, Nullability::NonNullable),
604                    Scalar::utf8("one", Nullability::Nullable),
605                ),
606                (
607                    Scalar::primitive(2i32, Nullability::NonNullable),
608                    Scalar::null(DType::Utf8(Nullability::Nullable)),
609                ),
610            ],
611        )?;
612
613        let datum = scalar.to_arrow_datum()?;
614        let (array, is_scalar) = datum.get();
615        assert!(is_scalar);
616        let map = array
617            .as_any()
618            .downcast_ref::<MapArray>()
619            .expect("map scalar should convert to MapArray");
620        assert_eq!(map.len(), 1);
621        assert_eq!(map.value_offsets(), &[0, 2]);
622        assert!(map.is_valid(0));
623        assert_eq!(
624            map.keys()
625                .as_any()
626                .downcast_ref::<Int32Array>()
627                .expect("map key array should be Int32")
628                .values(),
629            &[1, 2]
630        );
631        let values = map
632            .values()
633            .as_any()
634            .downcast_ref::<StringViewArray>()
635            .expect("map value array should be StringView");
636        assert_eq!(values.value(0), "one");
637        assert!(values.is_null(1));
638
639        Ok(())
640    }
641
642    #[test]
643    fn map_decimal_scalar_to_arrow_preserves_decimal_type() -> VortexResult<()> {
644        let decimal_dtype = DecimalDType::new(9, 2);
645        let dtype = DType::map(
646            DType::Primitive(PType::I32, Nullability::NonNullable),
647            DType::Decimal(decimal_dtype, Nullability::Nullable),
648            false,
649            Nullability::NonNullable,
650        )?;
651        let scalar = Scalar::try_map(
652            dtype,
653            [(
654                Scalar::primitive(1i32, Nullability::NonNullable),
655                Scalar::decimal(
656                    DecimalValue::I256(i256::from_i128(12345)),
657                    decimal_dtype,
658                    Nullability::Nullable,
659                ),
660            )],
661        )?;
662
663        let datum = scalar.to_arrow_datum()?;
664        let (array, is_scalar) = datum.get();
665        assert!(is_scalar);
666        let map = array
667            .as_any()
668            .downcast_ref::<MapArray>()
669            .expect("map scalar should convert to MapArray");
670        assert_eq!(map.values().data_type(), &DataType::Decimal128(9, 2));
671
672        let wide_decimal_dtype = DecimalDType::new(39, 2);
673        let dtype = DType::map(
674            DType::Primitive(PType::I32, Nullability::NonNullable),
675            DType::Decimal(wide_decimal_dtype, Nullability::Nullable),
676            false,
677            Nullability::NonNullable,
678        )?;
679        let scalar = Scalar::try_map(
680            dtype,
681            [(
682                Scalar::primitive(1i32, Nullability::NonNullable),
683                Scalar::decimal(
684                    DecimalValue::I64(12345),
685                    wide_decimal_dtype,
686                    Nullability::Nullable,
687                ),
688            )],
689        )?;
690
691        let datum = scalar.to_arrow_datum()?;
692        let (array, is_scalar) = datum.get();
693        assert!(is_scalar);
694        let map = array
695            .as_any()
696            .downcast_ref::<MapArray>()
697            .expect("map scalar should convert to MapArray");
698        assert_eq!(map.values().data_type(), &DataType::Decimal256(39, 2));
699
700        Ok(())
701    }
702
703    #[test]
704    fn map_scalar_with_unsupported_nested_value_errors_without_panic() -> VortexResult<()> {
705        let struct_dtype = DType::Struct(
706            StructFields::from_iter([(
707                "field1",
708                FieldDType::from(DType::Primitive(PType::I32, Nullability::NonNullable)),
709            )]),
710            Nullability::NonNullable,
711        );
712        let dtype = DType::map(
713            DType::Primitive(PType::I32, Nullability::NonNullable),
714            struct_dtype.clone(),
715            false,
716            Nullability::NonNullable,
717        )?;
718        let scalar = Scalar::try_map(
719            dtype,
720            [(
721                Scalar::primitive(1i32, Nullability::NonNullable),
722                Scalar::struct_(
723                    struct_dtype,
724                    vec![Scalar::primitive(42i32, Nullability::NonNullable)],
725                ),
726            )],
727        )?;
728
729        let error = scalar
730            .to_arrow_datum()
731            .err()
732            .expect("unsupported nested map value should error");
733        assert!(error.to_string().contains("struct scalar conversion"));
734
735        Ok(())
736    }
737
738    #[test]
739    fn test_struct_scalar_to_arrow_todo() {
740        let struct_dtype = DType::Struct(
741            StructFields::from_iter([(
742                "field1",
743                FieldDType::from(DType::Primitive(PType::I32, Nullability::NonNullable)),
744            )]),
745            Nullability::NonNullable,
746        );
747
748        let struct_scalar = Scalar::struct_(
749            struct_dtype,
750            vec![Scalar::primitive(42i32, Nullability::NonNullable)],
751        );
752        let error = struct_scalar
753            .to_arrow_datum()
754            .err()
755            .expect("struct scalar should error");
756        assert!(error.to_string().contains("struct scalar conversion"));
757    }
758
759    #[test]
760    fn test_list_scalar_to_arrow_todo() {
761        let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
762        let list_scalar = Scalar::list(
763            element_dtype,
764            vec![
765                Scalar::primitive(1i32, Nullability::NonNullable),
766                Scalar::primitive(2i32, Nullability::NonNullable),
767            ],
768            Nullability::NonNullable,
769        );
770
771        let error = list_scalar
772            .to_arrow_datum()
773            .err()
774            .expect("list scalar should error");
775        assert!(error.to_string().contains("list scalar conversion"));
776    }
777
778    #[test]
779    fn test_non_temporal_extension_to_arrow_todo() {
780        #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
781        struct SomeExt;
782        impl ExtVTable for SomeExt {
783            type Metadata = String;
784            type NativeValue<'a> = &'a str;
785
786            #[expect(clippy::disallowed_methods, reason = "test-only id")]
787            fn id(&self) -> ExtId {
788                ExtId::new("some_ext")
789            }
790
791            fn serialize_metadata(&self, _options: &Self::Metadata) -> VortexResult<Vec<u8>> {
792                vortex_bail!("not implemented")
793            }
794
795            fn deserialize_metadata(&self, _data: &[u8]) -> VortexResult<Self::Metadata> {
796                vortex_bail!("not implemented")
797            }
798
799            fn validate_dtype(_ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
800                Ok(())
801            }
802
803            fn unpack_native<'a>(
804                _ext_dtype: &'a ExtDType<Self>,
805                _storage_value: &'a ScalarValue,
806            ) -> VortexResult<Self::NativeValue<'a>> {
807                Ok("")
808            }
809        }
810
811        let scalar = Scalar::extension::<SomeExt>(
812            "".into(),
813            Scalar::primitive(42i32, Nullability::NonNullable),
814        );
815
816        let error = scalar
817            .to_arrow_datum()
818            .err()
819            .expect("non-temporal extension scalar should error");
820        assert!(
821            error
822                .to_string()
823                .contains("Cannot convert extension scalar")
824        );
825    }
826
827    #[rstest]
828    #[case(TimeUnit::Nanoseconds, PType::I64, 123456789i64)]
829    #[case(TimeUnit::Microseconds, PType::I64, 123456789i64)]
830    #[case(TimeUnit::Milliseconds, PType::I32, 123456i64)]
831    #[case(TimeUnit::Seconds, PType::I32, 1234i64)]
832    fn test_temporal_time_to_arrow(
833        #[case] time_unit: TimeUnit,
834        #[case] ptype: PType,
835        #[case] value: i64,
836    ) {
837        let scalar = Scalar::extension::<Time>(
838            time_unit,
839            match ptype {
840                PType::I32 => {
841                    let i32_value = i32::try_from(value).expect("test value should fit in i32");
842                    Scalar::primitive(i32_value, Nullability::NonNullable)
843                }
844                PType::I64 => Scalar::primitive(value, Nullability::NonNullable),
845                _ => unreachable!(),
846            },
847        );
848
849        let result = scalar.to_arrow_datum();
850        assert!(result.is_ok());
851    }
852
853    #[rstest]
854    #[case(TimeUnit::Milliseconds, PType::I64, 1234567890000i64)]
855    #[case(TimeUnit::Days, PType::I32, 19000i64)]
856    fn test_temporal_date_to_arrow(
857        #[case] time_unit: TimeUnit,
858        #[case] ptype: PType,
859        #[case] value: i64,
860    ) {
861        let scalar = Scalar::extension::<Date>(
862            time_unit,
863            match ptype {
864                PType::I32 => {
865                    let i32_value = i32::try_from(value).expect("test value should fit in i32");
866                    Scalar::primitive(i32_value, Nullability::NonNullable)
867                }
868                PType::I64 => Scalar::primitive(value, Nullability::NonNullable),
869                _ => unreachable!(),
870            },
871        );
872
873        let result = scalar.to_arrow_datum();
874        assert!(result.is_ok());
875    }
876
877    #[rstest]
878    #[case(TimeUnit::Nanoseconds, 1234567890000000000i64)]
879    #[case(TimeUnit::Microseconds, 1234567890000000i64)]
880    #[case(TimeUnit::Milliseconds, 1234567890000i64)]
881    #[case(TimeUnit::Seconds, 1234567890i64)]
882    fn test_temporal_timestamp_to_arrow(#[case] time_unit: TimeUnit, #[case] value: i64) {
883        let scalar = Scalar::extension::<Timestamp>(
884            TimestampOptions {
885                unit: time_unit,
886                tz: None,
887            },
888            Scalar::primitive(value, Nullability::NonNullable),
889        );
890
891        let result = scalar.to_arrow_datum();
892        assert!(result.is_ok());
893    }
894
895    #[rstest]
896    #[case(TimeUnit::Nanoseconds, "UTC", 1234567890000000000i64)]
897    #[case(TimeUnit::Microseconds, "America/New_York", 1234567890000000i64)]
898    #[case(TimeUnit::Microseconds, "Asia/Qatar", 1234567890000000i64)]
899    #[case(TimeUnit::Microseconds, "Australia/Sydney", 1234567890000000i64)]
900    #[case(TimeUnit::Milliseconds, "Pacific/Honolulu", 1234567890000i64)]
901    #[case(TimeUnit::Seconds, "GMT", 1234567890i64)]
902    fn test_temporal_timestamp_tz_to_arrow(
903        #[case] time_unit: TimeUnit,
904        #[case] tz: &str,
905        #[case] value: i64,
906    ) {
907        let scalar = Scalar::extension::<Timestamp>(
908            TimestampOptions {
909                unit: time_unit,
910                tz: Some(tz.into()),
911            },
912            Scalar::primitive(value, Nullability::NonNullable),
913        );
914
915        let result = scalar.to_arrow_datum();
916        assert!(result.is_ok());
917    }
918
919    #[test]
920    fn test_temporal_with_null_value() {
921        let scalar = Scalar::extension::<Time>(
922            TimeUnit::Milliseconds,
923            Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
924        );
925
926        let _result = scalar.to_arrow_datum().unwrap();
927    }
928
929    #[test]
930    #[should_panic(expected = "DType utf8 is not a primitive type")]
931    fn test_temporal_non_primitive_storage_error() {
932        let _scalar = Scalar::extension::<Time>(
933            TimeUnit::Nanoseconds,
934            Scalar::utf8("not a timestamp", Nullability::NonNullable),
935        );
936    }
937}
vortex_arrow/scalar.rs

vortex_arrow/
scalar.rs
