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::BinaryViewArray;
312    use arrow_array::BooleanArray;
313    use arrow_array::Date32Array;
314    use arrow_array::Date64Array;
315    use arrow_array::Decimal128Array;
316    use arrow_array::Decimal256Array;
317    use arrow_array::Float16Array;
318    use arrow_array::Float32Array;
319    use arrow_array::Float64Array;
320    use arrow_array::Int8Array;
321    use arrow_array::Int16Array;
322    use arrow_array::Int32Array;
323    use arrow_array::Int64Array;
324    use arrow_array::MapArray;
325    use arrow_array::NullArray;
326    use arrow_array::StringViewArray;
327    use arrow_array::Time32MillisecondArray;
328    use arrow_array::Time32SecondArray;
329    use arrow_array::Time64MicrosecondArray;
330    use arrow_array::Time64NanosecondArray;
331    use arrow_array::TimestampMicrosecondArray;
332    use arrow_array::TimestampMillisecondArray;
333    use arrow_array::TimestampNanosecondArray;
334    use arrow_array::TimestampSecondArray;
335    use arrow_array::UInt8Array;
336    use arrow_array::UInt16Array;
337    use arrow_array::UInt32Array;
338    use arrow_array::UInt64Array;
339    use arrow_schema::DataType;
340    use rstest::rstest;
341    use vortex_array::dtype::DType;
342    use vortex_array::dtype::DecimalDType;
343    use vortex_array::dtype::FieldDType;
344    use vortex_array::dtype::NativeDType;
345    use vortex_array::dtype::Nullability;
346    use vortex_array::dtype::PType;
347    use vortex_array::dtype::StructFields;
348    use vortex_array::dtype::extension::ExtDType;
349    use vortex_array::dtype::extension::ExtId;
350    use vortex_array::dtype::extension::ExtVTable;
351    use vortex_array::dtype::i256;
352    use vortex_array::extension::datetime::Date;
353    use vortex_array::extension::datetime::Time;
354    use vortex_array::extension::datetime::TimeUnit;
355    use vortex_array::extension::datetime::Timestamp;
356    use vortex_array::extension::datetime::TimestampOptions;
357    use vortex_array::scalar::DecimalValue;
358    use vortex_array::scalar::Scalar;
359    use vortex_array::scalar::ScalarValue;
360    use vortex_error::VortexResult;
361    use vortex_error::vortex_bail;
362
363    use super::ToArrowDatum;
364
365    #[test]
366    fn test_null_scalar_to_arrow() {
367        let scalar = Scalar::null(DType::Null);
368        let datum = scalar.to_arrow_datum().unwrap();
369        assert_eq!(datum.get().0.to_data(), NullArray::new(1).to_data());
370    }
371
372    #[test]
373    fn test_bool_scalar_to_arrow() {
374        let scalar = Scalar::bool(true, Nullability::NonNullable);
375        assert_arrow_scalar(&scalar, &BooleanArray::from(vec![true])).unwrap();
376    }
377
378    #[test]
379    fn test_null_bool_scalar_to_arrow() {
380        let scalar = Scalar::null(bool::dtype().as_nullable());
381        assert_arrow_scalar(&scalar, &BooleanArray::from(vec![None::<bool>])).unwrap();
382    }
383
384    #[test]
385    fn test_primitive_u8_to_arrow() {
386        let scalar = Scalar::primitive(42u8, Nullability::NonNullable);
387        assert_arrow_scalar(&scalar, &UInt8Array::from(vec![42u8])).unwrap();
388    }
389
390    #[test]
391    fn test_primitive_u16_to_arrow() {
392        let scalar = Scalar::primitive(1000u16, Nullability::NonNullable);
393        assert_arrow_scalar(&scalar, &UInt16Array::from(vec![1000u16])).unwrap();
394    }
395
396    #[test]
397    fn test_primitive_u32_to_arrow() {
398        let scalar = Scalar::primitive(100000u32, Nullability::NonNullable);
399        assert_arrow_scalar(&scalar, &UInt32Array::from(vec![100000u32])).unwrap();
400    }
401
402    #[test]
403    fn test_primitive_u64_to_arrow() {
404        let scalar = Scalar::primitive(10000000000u64, Nullability::NonNullable);
405        assert_arrow_scalar(&scalar, &UInt64Array::from(vec![10000000000u64])).unwrap();
406    }
407
408    #[test]
409    fn test_primitive_i8_to_arrow() {
410        let scalar = Scalar::primitive(-42i8, Nullability::NonNullable);
411        assert_arrow_scalar(&scalar, &Int8Array::from(vec![-42i8])).unwrap();
412    }
413
414    #[test]
415    fn test_primitive_i16_to_arrow() {
416        let scalar = Scalar::primitive(-1000i16, Nullability::NonNullable);
417        assert_arrow_scalar(&scalar, &Int16Array::from(vec![-1000i16])).unwrap();
418    }
419
420    #[test]
421    fn test_primitive_i32_to_arrow() {
422        let scalar = Scalar::primitive(-100000i32, Nullability::NonNullable);
423        assert_arrow_scalar(&scalar, &Int32Array::from(vec![-100000i32])).unwrap();
424    }
425
426    #[test]
427    fn test_primitive_i64_to_arrow() {
428        let scalar = Scalar::primitive(-10000000000i64, Nullability::NonNullable);
429        assert_arrow_scalar(&scalar, &Int64Array::from(vec![-10000000000i64])).unwrap();
430    }
431
432    #[test]
433    fn test_primitive_f16_to_arrow() {
434        use vortex_array::dtype::half::f16;
435
436        let scalar = Scalar::primitive(f16::from_f32(1.234), Nullability::NonNullable);
437        assert_arrow_scalar(&scalar, &Float16Array::from(vec![f16::from_f32(1.234)])).unwrap();
438    }
439
440    #[test]
441    fn test_primitive_f32_to_arrow() {
442        let scalar = Scalar::primitive(1.234f32, Nullability::NonNullable);
443        assert_arrow_scalar(&scalar, &Float32Array::from(vec![1.234f32])).unwrap();
444    }
445
446    #[test]
447    fn test_primitive_f64_to_arrow() {
448        let scalar = Scalar::primitive(1.234567890123f64, Nullability::NonNullable);
449        assert_arrow_scalar(&scalar, &Float64Array::from(vec![1.234567890123f64])).unwrap();
450    }
451
452    #[test]
453    fn test_null_primitive_to_arrow() {
454        let scalar = Scalar::null(i32::dtype().as_nullable());
455        assert_arrow_scalar(&scalar, &Int32Array::from(vec![None::<i32>])).unwrap();
456    }
457
458    #[test]
459    fn test_utf8_scalar_to_arrow() {
460        let scalar = Scalar::utf8("hello world".to_string(), Nullability::NonNullable);
461        assert_arrow_scalar(&scalar, &StringViewArray::from(vec!["hello world"])).unwrap();
462    }
463
464    #[test]
465    fn test_null_utf8_scalar_to_arrow() {
466        let scalar = Scalar::null(String::dtype().as_nullable());
467        assert_arrow_scalar(&scalar, &StringViewArray::from(vec![None::<&str>])).unwrap();
468    }
469
470    #[test]
471    fn test_binary_scalar_to_arrow() {
472        let data = vec![1u8, 2, 3, 4, 5];
473        let scalar = Scalar::binary(data, Nullability::NonNullable);
474        assert_arrow_scalar(
475            &scalar,
476            &BinaryViewArray::from(vec![&[1u8, 2, 3, 4, 5][..]]),
477        )
478        .unwrap();
479    }
480
481    #[test]
482    fn test_null_binary_scalar_to_arrow() {
483        let scalar = Scalar::null(DType::Binary(Nullability::Nullable));
484        assert_arrow_scalar(&scalar, &BinaryViewArray::from(vec![None::<&[u8]>])).unwrap();
485    }
486
487    fn assert_arrow_scalar(scalar: &Scalar, expected: &dyn Array) -> VortexResult<()> {
488        let datum = scalar.to_arrow_datum()?;
489        let (array, is_scalar) = datum.get();
490        assert!(is_scalar);
491        assert_eq!(array.to_data(), expected.to_data());
492        Ok(())
493    }
494
495    fn assert_arrow_scalar_data_type(
496        scalar: &Scalar,
497        expected: DataType,
498        value: Option<i128>,
499    ) -> VortexResult<()> {
500        let datum = scalar.to_arrow_datum()?;
501        let (array, is_scalar) = datum.get();
502        assert!(is_scalar);
503        assert_eq!(array.data_type(), &expected);
504        let expected: Arc<dyn Array> = match expected {
505            DataType::Decimal128(precision, scale) => Arc::new(
506                Decimal128Array::from(vec![value]).with_precision_and_scale(precision, scale)?,
507            ),
508            DataType::Decimal256(precision, scale) => Arc::new(
509                Decimal256Array::from(vec![value.map(arrow_buffer::i256::from_i128)])
510                    .with_precision_and_scale(precision, scale)?,
511            ),
512            _ => unreachable!(),
513        };
514        assert_eq!(array.to_data(), expected.to_data());
515        Ok(())
516    }
517
518    #[test]
519    fn test_decimal_scalars_to_arrow() -> VortexResult<()> {
520        // Test various decimal value types
521        let decimal_dtype = DecimalDType::new(5, 2);
522
523        let scalar_i8 = Scalar::decimal(
524            DecimalValue::I8(100),
525            decimal_dtype,
526            Nullability::NonNullable,
527        );
528        assert_arrow_scalar_data_type(&scalar_i8, DataType::Decimal128(5, 2), Some(100))?;
529
530        let scalar_i16 = Scalar::decimal(
531            DecimalValue::I16(10000),
532            decimal_dtype,
533            Nullability::NonNullable,
534        );
535        assert_arrow_scalar_data_type(&scalar_i16, DataType::Decimal128(5, 2), Some(10000))?;
536
537        let scalar_i32 = Scalar::decimal(
538            DecimalValue::I32(99999),
539            decimal_dtype,
540            Nullability::NonNullable,
541        );
542        assert_arrow_scalar_data_type(&scalar_i32, DataType::Decimal128(5, 2), Some(99999))?;
543
544        let scalar_i64 = Scalar::decimal(
545            DecimalValue::I64(99999),
546            decimal_dtype,
547            Nullability::NonNullable,
548        );
549        assert_arrow_scalar_data_type(&scalar_i64, DataType::Decimal128(5, 2), Some(99999))?;
550
551        let scalar_i128 = Scalar::decimal(
552            DecimalValue::I128(99999),
553            decimal_dtype,
554            Nullability::NonNullable,
555        );
556        assert_arrow_scalar_data_type(&scalar_i128, DataType::Decimal128(5, 2), Some(99999))?;
557
558        let value_i256 = i256::from_i128(99999);
559        let scalar_i256 = Scalar::decimal(
560            DecimalValue::I256(value_i256),
561            decimal_dtype,
562            Nullability::NonNullable,
563        );
564        assert_arrow_scalar_data_type(&scalar_i256, DataType::Decimal128(5, 2), Some(99999))?;
565
566        Ok(())
567    }
568
569    #[test]
570    fn decimal_i64_with_wide_precision_exports_decimal256() -> VortexResult<()> {
571        let scalar = Scalar::decimal(
572            DecimalValue::I64(1),
573            DecimalDType::new(39, 0),
574            Nullability::NonNullable,
575        );
576
577        assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 0), Some(1))
578    }
579
580    #[test]
581    fn decimal_i256_with_narrow_precision_exports_decimal128() -> VortexResult<()> {
582        let scalar = Scalar::decimal(
583            DecimalValue::I256(i256::from_i128(1234)),
584            DecimalDType::new(4, 2),
585            Nullability::NonNullable,
586        );
587
588        assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(4, 2), Some(1234))
589    }
590
591    #[test]
592    fn test_null_decimal_to_arrow() -> VortexResult<()> {
593        let decimal_dtype = DecimalDType::new(10, 2);
594        let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable));
595        assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(10, 2), None)?;
596
597        let decimal_dtype = DecimalDType::new(39, 2);
598        let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable));
599        assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 2), None)
600    }
601
602    #[test]
603    fn decimal_scalar_to_arrow_preserves_precision_and_scale() -> VortexResult<()> {
604        let decimal_dtype = DecimalDType::new(12, 3);
605        let scalar = Scalar::decimal(
606            DecimalValue::I128(12345),
607            decimal_dtype,
608            Nullability::NonNullable,
609        );
610
611        let datum = scalar.to_arrow_datum()?;
612        let (array, is_scalar) = datum.get();
613        assert!(is_scalar);
614        let decimal = array
615            .as_any()
616            .downcast_ref::<Decimal128Array>()
617            .expect("decimal scalar should convert to Decimal128");
618        assert_eq!(decimal.precision(), 12);
619        assert_eq!(decimal.scale(), 3);
620
621        assert_eq!(decimal.iter().collect::<Vec<_>>(), vec![Some(12345)]);
622        Ok(())
623    }
624
625    #[test]
626    fn test_map_scalar_to_arrow() -> VortexResult<()> {
627        let dtype = DType::map(
628            DType::Primitive(PType::I32, Nullability::NonNullable),
629            DType::Utf8(Nullability::Nullable),
630            true,
631            Nullability::Nullable,
632        )?;
633        let scalar = Scalar::try_map(
634            dtype,
635            [
636                (
637                    Scalar::primitive(1i32, Nullability::NonNullable),
638                    Scalar::utf8("one", Nullability::Nullable),
639                ),
640                (
641                    Scalar::primitive(2i32, Nullability::NonNullable),
642                    Scalar::null(DType::Utf8(Nullability::Nullable)),
643                ),
644            ],
645        )?;
646
647        let datum = scalar.to_arrow_datum()?;
648        let (array, is_scalar) = datum.get();
649        assert!(is_scalar);
650        let map = array
651            .as_any()
652            .downcast_ref::<MapArray>()
653            .expect("map scalar should convert to MapArray");
654        assert_eq!(map.len(), 1);
655        assert_eq!(map.value_offsets(), &[0, 2]);
656        assert!(map.is_valid(0));
657        assert_eq!(
658            map.keys()
659                .as_any()
660                .downcast_ref::<Int32Array>()
661                .expect("map key array should be Int32")
662                .values(),
663            &[1, 2]
664        );
665        let values = map
666            .values()
667            .as_any()
668            .downcast_ref::<StringViewArray>()
669            .expect("map value array should be StringView");
670        assert_eq!(values.value(0), "one");
671        assert!(values.is_null(1));
672
673        let DataType::Map(entries, true) = map.data_type() else {
674            panic!("map scalar must retain sorted keys");
675        };
676        assert!(!entries.is_nullable());
677        let DataType::Struct(fields) = entries.data_type() else {
678            unreachable!()
679        };
680        assert!(!fields[0].is_nullable());
681        assert!(fields[1].is_nullable());
682        Ok(())
683    }
684
685    #[test]
686    fn map_decimal_scalar_to_arrow_preserves_decimal_type() -> VortexResult<()> {
687        let decimal_dtype = DecimalDType::new(9, 2);
688        let dtype = DType::map(
689            DType::Primitive(PType::I32, Nullability::NonNullable),
690            DType::Decimal(decimal_dtype, Nullability::Nullable),
691            false,
692            Nullability::NonNullable,
693        )?;
694        let scalar = Scalar::try_map(
695            dtype,
696            [(
697                Scalar::primitive(1i32, Nullability::NonNullable),
698                Scalar::decimal(
699                    DecimalValue::I256(i256::from_i128(12345)),
700                    decimal_dtype,
701                    Nullability::Nullable,
702                ),
703            )],
704        )?;
705
706        let datum = scalar.to_arrow_datum()?;
707        let (array, is_scalar) = datum.get();
708        assert!(is_scalar);
709        let map = array
710            .as_any()
711            .downcast_ref::<MapArray>()
712            .expect("map scalar should convert to MapArray");
713        assert_eq!(map.values().data_type(), &DataType::Decimal128(9, 2));
714        assert_eq!(map.len(), 1);
715        assert!(map.is_valid(0));
716        assert_eq!(map.value_offsets(), &[0, 1]);
717        assert_eq!(map.keys().to_data(), Int32Array::from(vec![1]).to_data());
718        assert_eq!(
719            map.values().to_data(),
720            Decimal128Array::from(vec![Some(12345)])
721                .with_precision_and_scale(9, 2)?
722                .to_data()
723        );
724
725        let wide_decimal_dtype = DecimalDType::new(39, 2);
726        let dtype = DType::map(
727            DType::Primitive(PType::I32, Nullability::NonNullable),
728            DType::Decimal(wide_decimal_dtype, Nullability::Nullable),
729            false,
730            Nullability::NonNullable,
731        )?;
732        let scalar = Scalar::try_map(
733            dtype,
734            [(
735                Scalar::primitive(1i32, Nullability::NonNullable),
736                Scalar::decimal(
737                    DecimalValue::I64(12345),
738                    wide_decimal_dtype,
739                    Nullability::Nullable,
740                ),
741            )],
742        )?;
743
744        let datum = scalar.to_arrow_datum()?;
745        let (array, is_scalar) = datum.get();
746        assert!(is_scalar);
747        let map = array
748            .as_any()
749            .downcast_ref::<MapArray>()
750            .expect("map scalar should convert to MapArray");
751        assert_eq!(map.values().data_type(), &DataType::Decimal256(39, 2));
752        assert_eq!(map.len(), 1);
753        assert!(map.is_valid(0));
754        assert_eq!(map.value_offsets(), &[0, 1]);
755        assert_eq!(map.keys().to_data(), Int32Array::from(vec![1]).to_data());
756        assert_eq!(
757            map.values().to_data(),
758            Decimal256Array::from(vec![Some(arrow_buffer::i256::from_i128(12345))])
759                .with_precision_and_scale(39, 2)?
760                .to_data()
761        );
762
763        Ok(())
764    }
765
766    #[test]
767    fn map_scalar_with_unsupported_nested_value_errors_without_panic() -> VortexResult<()> {
768        let struct_dtype = DType::Struct(
769            StructFields::from_iter([(
770                "field1",
771                FieldDType::from(DType::Primitive(PType::I32, Nullability::NonNullable)),
772            )]),
773            Nullability::NonNullable,
774        );
775        let dtype = DType::map(
776            DType::Primitive(PType::I32, Nullability::NonNullable),
777            struct_dtype.clone(),
778            false,
779            Nullability::NonNullable,
780        )?;
781        let scalar = Scalar::try_map(
782            dtype,
783            [(
784                Scalar::primitive(1i32, Nullability::NonNullable),
785                Scalar::struct_(
786                    struct_dtype,
787                    vec![Scalar::primitive(42i32, Nullability::NonNullable)],
788                ),
789            )],
790        )?;
791
792        let error = scalar
793            .to_arrow_datum()
794            .err()
795            .expect("unsupported nested map value should error");
796        assert!(error.to_string().contains("struct scalar conversion"));
797
798        Ok(())
799    }
800
801    #[test]
802    fn test_struct_scalar_to_arrow_todo() {
803        let struct_dtype = DType::Struct(
804            StructFields::from_iter([(
805                "field1",
806                FieldDType::from(DType::Primitive(PType::I32, Nullability::NonNullable)),
807            )]),
808            Nullability::NonNullable,
809        );
810
811        let struct_scalar = Scalar::struct_(
812            struct_dtype,
813            vec![Scalar::primitive(42i32, Nullability::NonNullable)],
814        );
815        let error = struct_scalar
816            .to_arrow_datum()
817            .err()
818            .expect("struct scalar should error");
819        assert!(error.to_string().contains("struct scalar conversion"));
820    }
821
822    #[test]
823    fn test_list_scalar_to_arrow_todo() {
824        let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
825        let list_scalar = Scalar::list(
826            element_dtype,
827            vec![
828                Scalar::primitive(1i32, Nullability::NonNullable),
829                Scalar::primitive(2i32, Nullability::NonNullable),
830            ],
831            Nullability::NonNullable,
832        );
833
834        let error = list_scalar
835            .to_arrow_datum()
836            .err()
837            .expect("list scalar should error");
838        assert!(error.to_string().contains("list scalar conversion"));
839    }
840
841    #[test]
842    fn test_non_temporal_extension_to_arrow_todo() {
843        #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
844        struct SomeExt;
845        impl ExtVTable for SomeExt {
846            type Metadata = String;
847            type NativeValue<'a> = &'a str;
848
849            #[expect(clippy::disallowed_methods, reason = "test-only id")]
850            fn id(&self) -> ExtId {
851                ExtId::new("some_ext")
852            }
853
854            fn serialize_metadata(&self, _options: &Self::Metadata) -> VortexResult<Vec<u8>> {
855                vortex_bail!("not implemented")
856            }
857
858            fn deserialize_metadata(&self, _data: &[u8]) -> VortexResult<Self::Metadata> {
859                vortex_bail!("not implemented")
860            }
861
862            fn validate_dtype(_ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
863                Ok(())
864            }
865
866            fn unpack_native<'a>(
867                _ext_dtype: &'a ExtDType<Self>,
868                _storage_value: &'a ScalarValue,
869            ) -> VortexResult<Self::NativeValue<'a>> {
870                Ok("")
871            }
872        }
873
874        let scalar = Scalar::extension::<SomeExt>(
875            "".into(),
876            Scalar::primitive(42i32, Nullability::NonNullable),
877        );
878
879        let error = scalar
880            .to_arrow_datum()
881            .err()
882            .expect("non-temporal extension scalar should error");
883        assert!(
884            error
885                .to_string()
886                .contains("Cannot convert extension scalar")
887        );
888    }
889
890    #[rstest]
891    #[case(TimeUnit::Nanoseconds, PType::I64, 123456789i64)]
892    #[case(TimeUnit::Microseconds, PType::I64, 123456789i64)]
893    #[case(TimeUnit::Milliseconds, PType::I32, 123456i64)]
894    #[case(TimeUnit::Seconds, PType::I32, 1234i64)]
895    fn test_temporal_time_to_arrow(
896        #[case] time_unit: TimeUnit,
897        #[case] ptype: PType,
898        #[case] value: i64,
899    ) {
900        let scalar = Scalar::extension::<Time>(
901            time_unit,
902            match ptype {
903                PType::I32 => {
904                    let i32_value = i32::try_from(value).expect("test value should fit in i32");
905                    Scalar::primitive(i32_value, Nullability::NonNullable)
906                }
907                PType::I64 => Scalar::primitive(value, Nullability::NonNullable),
908                _ => unreachable!(),
909            },
910        );
911
912        match time_unit {
913            TimeUnit::Nanoseconds => {
914                assert_arrow_scalar(&scalar, &Time64NanosecondArray::from(vec![value]))
915            }
916            TimeUnit::Microseconds => {
917                assert_arrow_scalar(&scalar, &Time64MicrosecondArray::from(vec![value]))
918            }
919            TimeUnit::Milliseconds => assert_arrow_scalar(
920                &scalar,
921                &Time32MillisecondArray::from(vec![i32::try_from(value).unwrap()]),
922            ),
923            TimeUnit::Seconds => assert_arrow_scalar(
924                &scalar,
925                &Time32SecondArray::from(vec![i32::try_from(value).unwrap()]),
926            ),
927            TimeUnit::Days => unreachable!(),
928        }
929        .unwrap();
930    }
931
932    #[rstest]
933    #[case(TimeUnit::Milliseconds, PType::I64, 1234567890000i64)]
934    #[case(TimeUnit::Days, PType::I32, 19000i64)]
935    fn test_temporal_date_to_arrow(
936        #[case] time_unit: TimeUnit,
937        #[case] ptype: PType,
938        #[case] value: i64,
939    ) {
940        let scalar = Scalar::extension::<Date>(
941            time_unit,
942            match ptype {
943                PType::I32 => {
944                    let i32_value = i32::try_from(value).expect("test value should fit in i32");
945                    Scalar::primitive(i32_value, Nullability::NonNullable)
946                }
947                PType::I64 => Scalar::primitive(value, Nullability::NonNullable),
948                _ => unreachable!(),
949            },
950        );
951
952        match time_unit {
953            TimeUnit::Days => assert_arrow_scalar(
954                &scalar,
955                &Date32Array::from(vec![i32::try_from(value).unwrap()]),
956            ),
957            TimeUnit::Milliseconds => assert_arrow_scalar(&scalar, &Date64Array::from(vec![value])),
958            _ => unreachable!(),
959        }
960        .unwrap();
961    }
962
963    #[rstest]
964    #[case(TimeUnit::Nanoseconds, 1234567890000000000i64)]
965    #[case(TimeUnit::Microseconds, 1234567890000000i64)]
966    #[case(TimeUnit::Milliseconds, 1234567890000i64)]
967    #[case(TimeUnit::Seconds, 1234567890i64)]
968    fn test_temporal_timestamp_to_arrow(#[case] time_unit: TimeUnit, #[case] value: i64) {
969        let scalar = Scalar::extension::<Timestamp>(
970            TimestampOptions {
971                unit: time_unit,
972                tz: None,
973            },
974            Scalar::primitive(value, Nullability::NonNullable),
975        );
976
977        match time_unit {
978            TimeUnit::Seconds => {
979                assert_arrow_scalar(&scalar, &TimestampSecondArray::from(vec![value]))
980            }
981            TimeUnit::Milliseconds => {
982                assert_arrow_scalar(&scalar, &TimestampMillisecondArray::from(vec![value]))
983            }
984            TimeUnit::Microseconds => {
985                assert_arrow_scalar(&scalar, &TimestampMicrosecondArray::from(vec![value]))
986            }
987            TimeUnit::Nanoseconds => {
988                assert_arrow_scalar(&scalar, &TimestampNanosecondArray::from(vec![value]))
989            }
990            TimeUnit::Days => unreachable!(),
991        }
992        .unwrap();
993    }
994
995    #[rstest]
996    #[case(TimeUnit::Nanoseconds, "UTC", 1234567890000000000i64)]
997    #[case(TimeUnit::Microseconds, "America/New_York", 1234567890000000i64)]
998    #[case(TimeUnit::Microseconds, "Asia/Qatar", 1234567890000000i64)]
999    #[case(TimeUnit::Microseconds, "Australia/Sydney", 1234567890000000i64)]
1000    #[case(TimeUnit::Milliseconds, "Pacific/Honolulu", 1234567890000i64)]
1001    #[case(TimeUnit::Seconds, "GMT", 1234567890i64)]
1002    fn test_temporal_timestamp_tz_to_arrow(
1003        #[case] time_unit: TimeUnit,
1004        #[case] tz: &str,
1005        #[case] value: i64,
1006    ) {
1007        let scalar = Scalar::extension::<Timestamp>(
1008            TimestampOptions {
1009                unit: time_unit,
1010                tz: Some(tz.into()),
1011            },
1012            Scalar::primitive(value, Nullability::NonNullable),
1013        );
1014
1015        match time_unit {
1016            TimeUnit::Seconds => assert_arrow_scalar(
1017                &scalar,
1018                &TimestampSecondArray::from(vec![value]).with_timezone(tz),
1019            ),
1020            TimeUnit::Milliseconds => assert_arrow_scalar(
1021                &scalar,
1022                &TimestampMillisecondArray::from(vec![value]).with_timezone(tz),
1023            ),
1024            TimeUnit::Microseconds => assert_arrow_scalar(
1025                &scalar,
1026                &TimestampMicrosecondArray::from(vec![value]).with_timezone(tz),
1027            ),
1028            TimeUnit::Nanoseconds => assert_arrow_scalar(
1029                &scalar,
1030                &TimestampNanosecondArray::from(vec![value]).with_timezone(tz),
1031            ),
1032            TimeUnit::Days => unreachable!(),
1033        }
1034        .unwrap();
1035    }
1036
1037    #[test]
1038    fn test_temporal_with_null_value() {
1039        let scalar = Scalar::extension::<Time>(
1040            TimeUnit::Milliseconds,
1041            Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
1042        );
1043
1044        assert_arrow_scalar(&scalar, &Time32MillisecondArray::from(vec![None::<i32>])).unwrap();
1045    }
1046
1047    #[test]
1048    #[should_panic(expected = "DType utf8 is not a primitive type")]
1049    fn test_temporal_non_primitive_storage_error() {
1050        let _scalar = Scalar::extension::<Time>(
1051            TimeUnit::Nanoseconds,
1052            Scalar::utf8("not a timestamp", Nullability::NonNullable),
1053        );
1054    }
1055}
vortex_arrow/scalar.rs

vortex_arrow/
scalar.rs
