1use std::sync::Arc;
7
8use arrow_array::Scalar as ArrowScalar;
9use arrow_array::*;
10use vortex_array::dtype::DType;
11use vortex_array::dtype::PType;
12use vortex_array::extension::datetime::AnyTemporal;
13use vortex_array::extension::datetime::TemporalMetadata;
14use vortex_array::extension::datetime::TimeUnit;
15use vortex_array::scalar::BinaryScalar;
16use vortex_array::scalar::BoolScalar;
17use vortex_array::scalar::DecimalScalar;
18use vortex_array::scalar::DecimalValue;
19use vortex_array::scalar::ExtScalar;
20use vortex_array::scalar::PrimitiveScalar;
21use vortex_array::scalar::Scalar;
22use vortex_array::scalar::Utf8Scalar;
23use vortex_error::VortexError;
24use vortex_error::vortex_bail;
25use vortex_error::vortex_err;
26
27const SCALAR_ARRAY_LEN: usize = 1;
29
30macro_rules! value_to_arrow_scalar {
32 ($V:expr, $AR:ty) => {
33 Ok(std::sync::Arc::new(
34 $V.map(<$AR>::new_scalar)
35 .unwrap_or_else(|| arrow_array::Scalar::new(<$AR>::new_null(SCALAR_ARRAY_LEN))),
36 ))
37 };
38}
39
40macro_rules! timestamp_to_arrow_scalar {
42 ($V:expr, $TZ:expr, $AR:ty) => {{
43 let array = match $V {
44 Some(v) => <$AR>::new_scalar(v).into_inner(),
45 None => <$AR>::new_null(SCALAR_ARRAY_LEN),
46 }
47 .with_timezone_opt($TZ);
48 Ok(Arc::new(ArrowScalar::new(array)))
49 }};
50}
51
52pub trait ToArrowDatum {
57 fn to_arrow_datum(&self) -> Result<Arc<dyn Datum>, VortexError>;
59}
60
61impl ToArrowDatum for Scalar {
62 fn to_arrow_datum(&self) -> Result<Arc<dyn Datum>, VortexError> {
63 let value = self;
64 match value.dtype() {
65 DType::Null => Ok(Arc::new(NullArray::new(SCALAR_ARRAY_LEN))),
66 DType::Bool(_) => bool_to_arrow(value.as_bool()),
67 DType::Primitive(..) => primitive_to_arrow(value.as_primitive()),
68 DType::Decimal(..) => decimal_to_arrow(value.as_decimal()),
69 DType::Utf8(_) => utf8_to_arrow(value.as_utf8()),
70 DType::Binary(_) => binary_to_arrow(value.as_binary()),
71 DType::List(..) => unimplemented!("list scalar conversion"),
72 DType::FixedSizeList(..) => unimplemented!("fixed-size list scalar conversion"),
73 DType::Struct(..) => unimplemented!("struct scalar conversion"),
74 DType::Union(..) => unimplemented!("union scalar conversion"),
75 DType::Variant(_) => unimplemented!("Variant scalar conversion"),
76 DType::Extension(..) => extension_to_arrow(value.as_extension()),
77 }
78 }
79}
80
81fn bool_to_arrow(scalar: BoolScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
83 value_to_arrow_scalar!(scalar.value(), BooleanArray)
84}
85
86fn primitive_to_arrow(scalar: PrimitiveScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
88 match scalar.ptype() {
89 PType::U8 => value_to_arrow_scalar!(scalar.typed_value(), UInt8Array),
90 PType::U16 => value_to_arrow_scalar!(scalar.typed_value(), UInt16Array),
91 PType::U32 => value_to_arrow_scalar!(scalar.typed_value(), UInt32Array),
92 PType::U64 => value_to_arrow_scalar!(scalar.typed_value(), UInt64Array),
93 PType::I8 => value_to_arrow_scalar!(scalar.typed_value(), Int8Array),
94 PType::I16 => value_to_arrow_scalar!(scalar.typed_value(), Int16Array),
95 PType::I32 => value_to_arrow_scalar!(scalar.typed_value(), Int32Array),
96 PType::I64 => value_to_arrow_scalar!(scalar.typed_value(), Int64Array),
97 PType::F16 => value_to_arrow_scalar!(scalar.typed_value(), Float16Array),
98 PType::F32 => value_to_arrow_scalar!(scalar.typed_value(), Float32Array),
99 PType::F64 => value_to_arrow_scalar!(scalar.typed_value(), Float64Array),
100 }
101}
102
103fn decimal_to_arrow(scalar: DecimalScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
105 match scalar.decimal_value() {
107 Some(DecimalValue::I8(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
108 Some(DecimalValue::I16(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
109 Some(DecimalValue::I32(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
110 Some(DecimalValue::I64(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
111 Some(DecimalValue::I128(v128)) => Ok(Arc::new(Decimal128Array::new_scalar(v128))),
112 Some(DecimalValue::I256(v256)) => Ok(Arc::new(Decimal256Array::new_scalar(v256.into()))),
113 None => Ok(Arc::new(arrow_array::Scalar::new(
114 Decimal128Array::new_null(SCALAR_ARRAY_LEN),
115 ))),
116 }
117}
118
119fn utf8_to_arrow(scalar: Utf8Scalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
121 value_to_arrow_scalar!(scalar.value(), StringViewArray)
122}
123
124fn binary_to_arrow(scalar: BinaryScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
126 value_to_arrow_scalar!(scalar.value(), BinaryViewArray)
127}
128
129fn extension_to_arrow(scalar: ExtScalar<'_>) -> Result<Arc<dyn Datum>, VortexError> {
133 let ext_dtype = scalar.ext_dtype();
134 let Some(temporal) = ext_dtype.metadata_opt::<AnyTemporal>() else {
135 vortex_bail!(
136 "Cannot convert extension scalar {} to Arrow",
137 ext_dtype.id()
138 )
139 };
140
141 let storage_scalar = scalar.to_storage_scalar();
142 let primitive = storage_scalar
143 .as_primitive_opt()
144 .ok_or_else(|| vortex_err!("Expected primitive scalar"))?;
145
146 match temporal {
147 TemporalMetadata::Timestamp(unit, tz) => {
148 let value = primitive.as_::<i64>();
149 match unit {
150 TimeUnit::Nanoseconds => {
151 timestamp_to_arrow_scalar!(value, tz.clone(), TimestampNanosecondArray)
152 }
153 TimeUnit::Microseconds => {
154 timestamp_to_arrow_scalar!(value, tz.clone(), TimestampMicrosecondArray)
155 }
156 TimeUnit::Milliseconds => {
157 timestamp_to_arrow_scalar!(value, tz.clone(), TimestampMillisecondArray)
158 }
159 TimeUnit::Seconds => {
160 timestamp_to_arrow_scalar!(value, tz.clone(), TimestampSecondArray)
161 }
162 TimeUnit::Days => {
163 vortex_bail!("Unsupported TimeUnit {unit} for {}", ext_dtype.id())
164 }
165 }
166 }
167 TemporalMetadata::Date(unit) => match unit {
168 TimeUnit::Milliseconds => {
169 value_to_arrow_scalar!(primitive.as_::<i64>(), Date64Array)
170 }
171 TimeUnit::Days => {
172 value_to_arrow_scalar!(primitive.as_::<i32>(), Date32Array)
173 }
174 TimeUnit::Nanoseconds | TimeUnit::Microseconds | TimeUnit::Seconds => {
175 vortex_bail!("Unsupported TimeUnit {unit} for {}", ext_dtype.id())
176 }
177 },
178 TemporalMetadata::Time(unit) => match unit {
179 TimeUnit::Nanoseconds => {
180 value_to_arrow_scalar!(primitive.as_::<i64>(), Time64NanosecondArray)
181 }
182 TimeUnit::Microseconds => {
183 value_to_arrow_scalar!(primitive.as_::<i64>(), Time64MicrosecondArray)
184 }
185 TimeUnit::Milliseconds => {
186 value_to_arrow_scalar!(primitive.as_::<i32>(), Time32MillisecondArray)
187 }
188 TimeUnit::Seconds => {
189 value_to_arrow_scalar!(primitive.as_::<i32>(), Time32SecondArray)
190 }
191 TimeUnit::Days => {
192 vortex_bail!("Unsupported TimeUnit {unit} for {}", ext_dtype.id())
193 }
194 },
195 }
196}
197
198#[cfg(test)]
199mod tests {
200 use std::sync::Arc;
201
202 use rstest::rstest;
203 use vortex_array::dtype::DType;
204 use vortex_array::dtype::DecimalDType;
205 use vortex_array::dtype::FieldDType;
206 use vortex_array::dtype::NativeDType;
207 use vortex_array::dtype::Nullability;
208 use vortex_array::dtype::PType;
209 use vortex_array::dtype::StructFields;
210 use vortex_array::dtype::extension::ExtDType;
211 use vortex_array::dtype::extension::ExtId;
212 use vortex_array::dtype::extension::ExtVTable;
213 use vortex_array::dtype::i256;
214 use vortex_array::extension::datetime::Date;
215 use vortex_array::extension::datetime::Time;
216 use vortex_array::extension::datetime::TimeUnit;
217 use vortex_array::extension::datetime::Timestamp;
218 use vortex_array::extension::datetime::TimestampOptions;
219 use vortex_array::scalar::DecimalValue;
220 use vortex_array::scalar::Scalar;
221 use vortex_array::scalar::ScalarValue;
222 use vortex_error::VortexResult;
223 use vortex_error::vortex_bail;
224
225 use super::ToArrowDatum;
226
227 #[test]
228 fn test_null_scalar_to_arrow() {
229 let scalar = Scalar::null(DType::Null);
230 let result = scalar.to_arrow_datum();
231 assert!(result.is_ok());
232 }
233
234 #[test]
235 fn test_bool_scalar_to_arrow() {
236 let scalar = Scalar::bool(true, Nullability::NonNullable);
237 let result = scalar.to_arrow_datum();
238 assert!(result.is_ok());
239 }
240
241 #[test]
242 fn test_null_bool_scalar_to_arrow() {
243 let scalar = Scalar::null(bool::dtype().as_nullable());
244 let result = scalar.to_arrow_datum();
245 assert!(result.is_ok());
246 }
247
248 #[test]
249 fn test_primitive_u8_to_arrow() {
250 let scalar = Scalar::primitive(42u8, Nullability::NonNullable);
251 let result = scalar.to_arrow_datum();
252 assert!(result.is_ok());
253 }
254
255 #[test]
256 fn test_primitive_u16_to_arrow() {
257 let scalar = Scalar::primitive(1000u16, Nullability::NonNullable);
258 let result = scalar.to_arrow_datum();
259 assert!(result.is_ok());
260 }
261
262 #[test]
263 fn test_primitive_u32_to_arrow() {
264 let scalar = Scalar::primitive(100000u32, Nullability::NonNullable);
265 let result = scalar.to_arrow_datum();
266 assert!(result.is_ok());
267 }
268
269 #[test]
270 fn test_primitive_u64_to_arrow() {
271 let scalar = Scalar::primitive(10000000000u64, Nullability::NonNullable);
272 let result = scalar.to_arrow_datum();
273 assert!(result.is_ok());
274 }
275
276 #[test]
277 fn test_primitive_i8_to_arrow() {
278 let scalar = Scalar::primitive(-42i8, Nullability::NonNullable);
279 let result = scalar.to_arrow_datum();
280 assert!(result.is_ok());
281 }
282
283 #[test]
284 fn test_primitive_i16_to_arrow() {
285 let scalar = Scalar::primitive(-1000i16, Nullability::NonNullable);
286 let result = scalar.to_arrow_datum();
287 assert!(result.is_ok());
288 }
289
290 #[test]
291 fn test_primitive_i32_to_arrow() {
292 let scalar = Scalar::primitive(-100000i32, Nullability::NonNullable);
293 let result = scalar.to_arrow_datum();
294 assert!(result.is_ok());
295 }
296
297 #[test]
298 fn test_primitive_i64_to_arrow() {
299 let scalar = Scalar::primitive(-10000000000i64, Nullability::NonNullable);
300 let result = scalar.to_arrow_datum();
301 assert!(result.is_ok());
302 }
303
304 #[test]
305 fn test_primitive_f16_to_arrow() {
306 use vortex_array::dtype::half::f16;
307
308 let scalar = Scalar::primitive(f16::from_f32(1.234), Nullability::NonNullable);
309 let result = scalar.to_arrow_datum();
310 assert!(result.is_ok());
311 }
312
313 #[test]
314 fn test_primitive_f32_to_arrow() {
315 let scalar = Scalar::primitive(1.234f32, Nullability::NonNullable);
316 let result = scalar.to_arrow_datum();
317 assert!(result.is_ok());
318 }
319
320 #[test]
321 fn test_primitive_f64_to_arrow() {
322 let scalar = Scalar::primitive(1.234567890123f64, Nullability::NonNullable);
323 let result = scalar.to_arrow_datum();
324 assert!(result.is_ok());
325 }
326
327 #[test]
328 fn test_null_primitive_to_arrow() {
329 let scalar = Scalar::null(i32::dtype().as_nullable());
330 let result = scalar.to_arrow_datum();
331 assert!(result.is_ok());
332 }
333
334 #[test]
335 fn test_utf8_scalar_to_arrow() {
336 let scalar = Scalar::utf8("hello world".to_string(), Nullability::NonNullable);
337 let result = scalar.to_arrow_datum();
338 assert!(result.is_ok());
339 }
340
341 #[test]
342 fn test_null_utf8_scalar_to_arrow() {
343 let scalar = Scalar::null(String::dtype().as_nullable());
344 let result = scalar.to_arrow_datum();
345 assert!(result.is_ok());
346 }
347
348 #[test]
349 fn test_binary_scalar_to_arrow() {
350 let data = vec![1u8, 2, 3, 4, 5];
351 let scalar = Scalar::binary(data, Nullability::NonNullable);
352 let result = scalar.to_arrow_datum();
353 assert!(result.is_ok());
354 }
355
356 #[test]
357 fn test_null_binary_scalar_to_arrow() {
358 let scalar = Scalar::null(DType::Binary(Nullability::Nullable));
359 let result = scalar.to_arrow_datum();
360 assert!(result.is_ok());
361 }
362
363 #[test]
364 fn test_decimal_scalars_to_arrow() {
365 let decimal_dtype = DecimalDType::new(5, 2);
367
368 let scalar_i8 = Scalar::decimal(
369 DecimalValue::I8(100),
370 decimal_dtype,
371 Nullability::NonNullable,
372 );
373 assert!(scalar_i8.to_arrow_datum().is_ok());
374
375 let scalar_i16 = Scalar::decimal(
376 DecimalValue::I16(10000),
377 decimal_dtype,
378 Nullability::NonNullable,
379 );
380 assert!(scalar_i16.to_arrow_datum().is_ok());
381
382 let scalar_i32 = Scalar::decimal(
383 DecimalValue::I32(99999),
384 decimal_dtype,
385 Nullability::NonNullable,
386 );
387 assert!(scalar_i32.to_arrow_datum().is_ok());
388
389 let scalar_i64 = Scalar::decimal(
390 DecimalValue::I64(99999),
391 decimal_dtype,
392 Nullability::NonNullable,
393 );
394 assert!(scalar_i64.to_arrow_datum().is_ok());
395
396 let scalar_i128 = Scalar::decimal(
397 DecimalValue::I128(99999),
398 decimal_dtype,
399 Nullability::NonNullable,
400 );
401 assert!(scalar_i128.to_arrow_datum().is_ok());
402
403 let value_i256 = i256::from_i128(99999);
406 let scalar_i256 = Scalar::decimal(
407 DecimalValue::I256(value_i256),
408 decimal_dtype,
409 Nullability::NonNullable,
410 );
411 assert!(scalar_i256.to_arrow_datum().is_ok());
412 }
413
414 #[test]
415 fn test_null_decimal_to_arrow() {
416 let decimal_dtype = DecimalDType::new(10, 2);
417 let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable));
418 let result = scalar.to_arrow_datum();
419 assert!(result.is_ok());
420 }
421
422 #[test]
423 #[should_panic(expected = "struct scalar conversion")]
424 fn test_struct_scalar_to_arrow_todo() {
425 let struct_dtype = DType::Struct(
426 StructFields::from_iter([(
427 "field1",
428 FieldDType::from(DType::Primitive(PType::I32, Nullability::NonNullable)),
429 )]),
430 Nullability::NonNullable,
431 );
432
433 let struct_scalar = Scalar::struct_(
434 struct_dtype,
435 vec![Scalar::primitive(42i32, Nullability::NonNullable)],
436 );
437 struct_scalar.to_arrow_datum().unwrap();
438 }
439
440 #[test]
441 #[should_panic(expected = "list scalar conversion")]
442 fn test_list_scalar_to_arrow_todo() {
443 let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
444 let list_scalar = Scalar::list(
445 element_dtype,
446 vec![
447 Scalar::primitive(1i32, Nullability::NonNullable),
448 Scalar::primitive(2i32, Nullability::NonNullable),
449 ],
450 Nullability::NonNullable,
451 );
452
453 list_scalar.to_arrow_datum().unwrap();
454 }
455
456 #[test]
457 #[should_panic(expected = "Cannot convert extension scalar")]
458 fn test_non_temporal_extension_to_arrow_todo() {
459 #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
460 struct SomeExt;
461 impl ExtVTable for SomeExt {
462 type Metadata = String;
463 type NativeValue<'a> = &'a str;
464
465 #[expect(clippy::disallowed_methods, reason = "test-only id")]
466 fn id(&self) -> ExtId {
467 ExtId::new("some_ext")
468 }
469
470 fn serialize_metadata(&self, _options: &Self::Metadata) -> VortexResult<Vec<u8>> {
471 vortex_bail!("not implemented")
472 }
473
474 fn deserialize_metadata(&self, _data: &[u8]) -> VortexResult<Self::Metadata> {
475 vortex_bail!("not implemented")
476 }
477
478 fn validate_dtype(_ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
479 Ok(())
480 }
481
482 fn unpack_native<'a>(
483 _ext_dtype: &'a ExtDType<Self>,
484 _storage_value: &'a ScalarValue,
485 ) -> VortexResult<Self::NativeValue<'a>> {
486 Ok("")
487 }
488 }
489
490 let scalar = Scalar::extension::<SomeExt>(
491 "".into(),
492 Scalar::primitive(42i32, Nullability::NonNullable),
493 );
494
495 scalar.to_arrow_datum().unwrap();
496 }
497
498 #[rstest]
499 #[case(TimeUnit::Nanoseconds, PType::I64, 123456789i64)]
500 #[case(TimeUnit::Microseconds, PType::I64, 123456789i64)]
501 #[case(TimeUnit::Milliseconds, PType::I32, 123456i64)]
502 #[case(TimeUnit::Seconds, PType::I32, 1234i64)]
503 fn test_temporal_time_to_arrow(
504 #[case] time_unit: TimeUnit,
505 #[case] ptype: PType,
506 #[case] value: i64,
507 ) {
508 let scalar = Scalar::extension::<Time>(
509 time_unit,
510 match ptype {
511 PType::I32 => {
512 let i32_value = i32::try_from(value).expect("test value should fit in i32");
513 Scalar::primitive(i32_value, Nullability::NonNullable)
514 }
515 PType::I64 => Scalar::primitive(value, Nullability::NonNullable),
516 _ => unreachable!(),
517 },
518 );
519
520 let result = scalar.to_arrow_datum();
521 assert!(result.is_ok());
522 }
523
524 #[rstest]
525 #[case(TimeUnit::Milliseconds, PType::I64, 1234567890000i64)]
526 #[case(TimeUnit::Days, PType::I32, 19000i64)]
527 fn test_temporal_date_to_arrow(
528 #[case] time_unit: TimeUnit,
529 #[case] ptype: PType,
530 #[case] value: i64,
531 ) {
532 let scalar = Scalar::extension::<Date>(
533 time_unit,
534 match ptype {
535 PType::I32 => {
536 let i32_value = i32::try_from(value).expect("test value should fit in i32");
537 Scalar::primitive(i32_value, Nullability::NonNullable)
538 }
539 PType::I64 => Scalar::primitive(value, Nullability::NonNullable),
540 _ => unreachable!(),
541 },
542 );
543
544 let result = scalar.to_arrow_datum();
545 assert!(result.is_ok());
546 }
547
548 #[rstest]
549 #[case(TimeUnit::Nanoseconds, 1234567890000000000i64)]
550 #[case(TimeUnit::Microseconds, 1234567890000000i64)]
551 #[case(TimeUnit::Milliseconds, 1234567890000i64)]
552 #[case(TimeUnit::Seconds, 1234567890i64)]
553 fn test_temporal_timestamp_to_arrow(#[case] time_unit: TimeUnit, #[case] value: i64) {
554 let scalar = Scalar::extension::<Timestamp>(
555 TimestampOptions {
556 unit: time_unit,
557 tz: None,
558 },
559 Scalar::primitive(value, Nullability::NonNullable),
560 );
561
562 let result = scalar.to_arrow_datum();
563 assert!(result.is_ok());
564 }
565
566 #[rstest]
567 #[case(TimeUnit::Nanoseconds, "UTC", 1234567890000000000i64)]
568 #[case(TimeUnit::Microseconds, "America/New_York", 1234567890000000i64)]
569 #[case(TimeUnit::Microseconds, "Asia/Qatar", 1234567890000000i64)]
570 #[case(TimeUnit::Microseconds, "Australia/Sydney", 1234567890000000i64)]
571 #[case(TimeUnit::Milliseconds, "Pacific/Honolulu", 1234567890000i64)]
572 #[case(TimeUnit::Seconds, "GMT", 1234567890i64)]
573 fn test_temporal_timestamp_tz_to_arrow(
574 #[case] time_unit: TimeUnit,
575 #[case] tz: &str,
576 #[case] value: i64,
577 ) {
578 let scalar = Scalar::extension::<Timestamp>(
579 TimestampOptions {
580 unit: time_unit,
581 tz: Some(tz.into()),
582 },
583 Scalar::primitive(value, Nullability::NonNullable),
584 );
585
586 let result = scalar.to_arrow_datum();
587 assert!(result.is_ok());
588 }
589
590 #[test]
591 fn test_temporal_with_null_value() {
592 let scalar = Scalar::extension::<Time>(
593 TimeUnit::Milliseconds,
594 Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
595 );
596
597 let _result = scalar.to_arrow_datum().unwrap();
598 }
599
600 #[test]
601 #[should_panic(expected = "DType utf8 is not a primitive type")]
602 fn test_temporal_non_primitive_storage_error() {
603 let _scalar = Scalar::extension::<Time>(
604 TimeUnit::Nanoseconds,
605 Scalar::utf8("not a timestamp", Nullability::NonNullable),
606 );
607 }
608}