1use std::cmp::Ordering;
7use std::hash::Hash;
8use std::hash::Hasher;
9
10use vortex_error::VortexResult;
11use vortex_error::vortex_ensure_eq;
12use vortex_error::vortex_panic;
13
14use crate::dtype::DType;
15use crate::dtype::NativeDType;
16use crate::dtype::PType;
17use crate::dtype::StructFields;
18use crate::scalar::Scalar;
19use crate::scalar::ScalarValue;
20
21impl Scalar {
22 pub fn null(dtype: DType) -> Self {
30 assert!(
31 dtype.is_nullable(),
32 "Cannot create null scalar with non-nullable dtype {dtype}"
33 );
34
35 Self { dtype, value: None }
36 }
37
38 pub fn null_native<T: NativeDType>() -> Self {
43 Self {
44 dtype: T::dtype().as_nullable(),
45 value: None,
46 }
47 }
48
49 #[cfg(test)]
59 pub fn new(dtype: DType, value: Option<ScalarValue>) -> Self {
60 use vortex_error::VortexExpect;
61
62 Self::try_new(dtype, value).vortex_expect("Failed to create Scalar")
63 }
64
65 pub fn try_new(dtype: DType, value: Option<ScalarValue>) -> VortexResult<Self> {
72 Self::validate(&dtype, value.as_ref())?;
73
74 Ok(Self { dtype, value })
75 }
76
77 pub unsafe fn new_unchecked(dtype: DType, value: Option<ScalarValue>) -> Self {
85 #[cfg(debug_assertions)]
86 {
87 use vortex_error::VortexExpect;
88
89 Self::validate(&dtype, value.as_ref())
90 .vortex_expect("Scalar::new_unchecked called with incompatible dtype and value");
91 }
92
93 Self { dtype, value }
94 }
95
96 pub fn default_value(dtype: &DType) -> Self {
115 Self::try_default_value(dtype)
116 .unwrap_or_else(|| vortex_panic!("{dtype} has no default value"))
117 }
118
119 pub(crate) fn try_default_value(dtype: &DType) -> Option<Self> {
121 let value = ScalarValue::try_default_value(dtype)?;
122 Self::try_new(dtype.clone(), value).ok()
123 }
124
125 pub fn zero_value(dtype: &DType) -> Self {
153 let value = ScalarValue::zero_value(dtype);
154
155 unsafe { Self::new_unchecked(dtype.clone(), Some(value)) }
157 }
158
159 pub fn eq_ignore_nullability(&self, other: &Self) -> bool {
163 self.dtype.eq_ignore_nullability(&other.dtype) && self.value == other.value
164 }
165
166 pub fn into_parts(self) -> (DType, Option<ScalarValue>) {
168 (self.dtype, self.value)
169 }
170
171 pub fn dtype(&self) -> &DType {
173 &self.dtype
174 }
175
176 pub fn value(&self) -> Option<&ScalarValue> {
178 self.value.as_ref()
179 }
180
181 pub fn into_value(self) -> Option<ScalarValue> {
184 self.value
185 }
186
187 pub fn is_valid(&self) -> bool {
189 self.value.is_some()
190 }
191
192 pub fn is_null(&self) -> bool {
194 self.value.is_none()
195 }
196
197 pub fn is_zero(&self) -> Option<bool> {
206 let value = self.value()?;
207
208 let is_zero = match self.dtype() {
209 DType::Null => vortex_panic!("non-null value somehow had `DType::Null`"),
210 DType::Bool(_) => !value.as_bool(),
211 DType::Primitive(..) => value.as_primitive().is_zero(),
212 DType::Decimal(..) => value.as_decimal().is_zero(),
213 DType::Utf8(_) => value.as_utf8().is_empty(),
214 DType::Binary(_) => value.as_binary().is_empty(),
215 DType::List(..) => value.as_list().is_empty(),
216 DType::Map(..) => self.as_map().is_empty(),
217 DType::FixedSizeList(_, list_size, _) => {
220 let list = self.as_list();
221 list.len() == *list_size as usize
222 && (0..list.len())
223 .all(|i| list.element(i).is_some_and(|e| e.is_zero() == Some(true)))
224 }
225 DType::Struct(..) => self
227 .as_struct()
228 .fields_iter()
229 .is_some_and(|mut fields| fields.all(|f| f.is_zero() == Some(true))),
230 DType::Union(..) => {
233 let union = self.as_union();
234 if union.child_index() != Some(0) {
235 false
236 } else {
237 union.child().and_then(|child| child.is_zero())?
238 }
239 }
240 DType::Variant(_) => self.as_variant().is_zero()?,
241 DType::Extension(_) => self.as_extension().to_storage_scalar().is_zero()?,
242 };
243
244 Some(is_zero)
245 }
246
247 pub fn primitive_reinterpret_cast(&self, ptype: PType) -> VortexResult<Self> {
253 let primitive = self.as_primitive();
254 if primitive.ptype() == ptype {
255 return Ok(self.clone());
256 }
257
258 vortex_ensure_eq!(
259 primitive.ptype().byte_width(),
260 ptype.byte_width(),
261 "can't reinterpret cast between integers of two different widths"
262 );
263
264 Scalar::try_new(
265 DType::Primitive(ptype, self.dtype().nullability()),
266 primitive
267 .pvalue()
268 .map(|p| p.reinterpret_cast(ptype))
269 .map(ScalarValue::Primitive),
270 )
271 }
272
273 pub fn approx_nbytes(&self) -> usize {
278 use crate::dtype::NativeDecimalType;
279 use crate::dtype::i256;
280
281 match self.dtype() {
282 DType::Null => 0,
283 DType::Bool(_) => 1,
284 DType::Primitive(ptype, _) => ptype.byte_width(),
285 DType::Decimal(dt, _) => {
286 if dt.precision() <= i128::MAX_PRECISION {
287 size_of::<i128>()
288 } else {
289 size_of::<i256>()
290 }
291 }
292 DType::Utf8(_) => self
293 .value()
294 .map_or_else(|| 0, |value| value.as_utf8().len()),
295 DType::Binary(_) => self
296 .value()
297 .map_or_else(|| 0, |value| value.as_binary().len()),
298 DType::List(..) | DType::FixedSizeList(..) => self
299 .as_list()
300 .elements()
301 .map(|fields| fields.into_iter().map(|f| f.approx_nbytes()).sum::<usize>())
302 .unwrap_or_default(),
303 DType::Map(..) => self
304 .as_map()
305 .entries()
306 .map(|(key, value)| key.approx_nbytes() + value.approx_nbytes())
307 .sum(),
308 DType::Struct(..) => self
309 .as_struct()
310 .fields_iter()
311 .map(|fields| fields.into_iter().map(|f| f.approx_nbytes()).sum::<usize>())
312 .unwrap_or_default(),
313 DType::Union(..) => self
314 .as_union()
315 .child()
316 .map_or(0, |value| 1 + value.approx_nbytes()),
317 DType::Variant(_) => self.as_variant().value().map_or(0, Scalar::approx_nbytes),
318 DType::Extension(_) => self.as_extension().to_storage_scalar().approx_nbytes(),
319 }
320 }
321}
322
323impl Hash for Scalar {
327 fn hash<H: Hasher>(&self, state: &mut H) {
328 self.dtype.hash_ignore_nullability(state);
329 self.value.hash(state);
330 }
331}
332
333impl PartialEq for Scalar {
340 fn eq(&self, other: &Self) -> bool {
341 self.dtype.eq_ignore_nullability(&other.dtype) && self.value == other.value
342 }
343}
344
345impl PartialOrd for Scalar {
346 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
382 if !self.dtype().eq_ignore_nullability(other.dtype()) {
383 return None;
384 }
385
386 partial_cmp_scalar_values(self.dtype(), self.value(), other.value())
387 }
388}
389
390fn partial_cmp_scalar_values(
392 dtype: &DType,
393 lhs: Option<&ScalarValue>,
394 rhs: Option<&ScalarValue>,
395) -> Option<Ordering> {
396 match (lhs, rhs) {
397 (None, None) => Some(Ordering::Equal),
398 (None, Some(_)) => Some(Ordering::Less),
399 (Some(_), None) => Some(Ordering::Greater),
400 (Some(lhs), Some(rhs)) => partial_cmp_non_null_scalar_values(dtype, lhs, rhs),
401 }
402}
403
404fn partial_cmp_non_null_scalar_values(
406 dtype: &DType,
407 lhs: &ScalarValue,
408 rhs: &ScalarValue,
409) -> Option<Ordering> {
410 match (lhs, rhs) {
413 (ScalarValue::Bool(lhs), ScalarValue::Bool(rhs)) => lhs.partial_cmp(rhs),
414 (ScalarValue::Primitive(lhs), ScalarValue::Primitive(rhs)) => lhs.partial_cmp(rhs),
415 (ScalarValue::Decimal(lhs), ScalarValue::Decimal(rhs)) => lhs.partial_cmp(rhs),
416 (ScalarValue::Utf8(lhs), ScalarValue::Utf8(rhs)) => lhs.partial_cmp(rhs),
417 (ScalarValue::Binary(lhs), ScalarValue::Binary(rhs)) => lhs.partial_cmp(rhs),
418 (ScalarValue::Tuple(lhs), ScalarValue::Tuple(rhs)) => {
421 partial_cmp_tuple_values(dtype, lhs, rhs)
422 }
423 (ScalarValue::Union(lhs), ScalarValue::Union(rhs)) => {
424 if lhs.type_id() != rhs.type_id() {
425 return None;
426 }
427
428 let DType::Union(variants, _) = dtype else {
429 return None;
430 };
431 let child_index = variants.tag_to_child_index(lhs.type_id())?;
432 let child_dtype = variants.variant_by_index(child_index)?;
433
434 partial_cmp_scalar_values(&child_dtype, lhs.child_value(), rhs.child_value())
435 }
436 (ScalarValue::Variant(_), ScalarValue::Variant(_)) => None,
439 _ => None,
440 }
441}
442
443fn partial_cmp_tuple_values(
445 dtype: &DType,
446 lhs: &[Option<ScalarValue>],
447 rhs: &[Option<ScalarValue>],
448) -> Option<Ordering> {
449 match dtype {
450 DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => {
451 partial_cmp_list_values(element_dtype, lhs, rhs)
452 }
453 DType::Struct(fields, _) => partial_cmp_struct_values(fields, lhs, rhs),
454 DType::Map(..) => None,
455 DType::Extension(ext_dtype) => {
456 partial_cmp_tuple_values(ext_dtype.storage_dtype(), lhs, rhs)
457 }
458 _ => None,
459 }
460}
461
462fn partial_cmp_list_values(
464 element_dtype: &DType,
465 lhs: &[Option<ScalarValue>],
466 rhs: &[Option<ScalarValue>],
467) -> Option<Ordering> {
468 for (lhs, rhs) in lhs.iter().zip(rhs.iter()) {
469 match partial_cmp_scalar_values(element_dtype, lhs.as_ref(), rhs.as_ref())? {
470 Ordering::Equal => continue,
471 ordering => return Some(ordering),
472 }
473 }
474
475 Some(lhs.len().cmp(&rhs.len()))
476}
477
478fn partial_cmp_struct_values(
480 fields: &StructFields,
481 lhs: &[Option<ScalarValue>],
482 rhs: &[Option<ScalarValue>],
483) -> Option<Ordering> {
484 if lhs.len() != fields.nfields() || rhs.len() != fields.nfields() {
485 return None;
486 }
487
488 for ((field_dtype, lhs), rhs) in fields.fields().zip(lhs.iter()).zip(rhs.iter()) {
489 match partial_cmp_scalar_values(&field_dtype, lhs.as_ref(), rhs.as_ref())? {
490 Ordering::Equal => continue,
491 ordering => return Some(ordering),
492 }
493 }
494
495 Some(Ordering::Equal)
496}
497
498#[cfg(test)]
499mod tests {
500 use std::sync::Arc;
501
502 use rstest::rstest;
503
504 use crate::dtype::DType;
505 use crate::dtype::Nullability;
506 use crate::dtype::PType;
507 use crate::dtype::StructFields;
508 use crate::scalar::Scalar;
509
510 fn i32_scalar(value: i32) -> Scalar {
511 Scalar::primitive::<i32>(value, Nullability::NonNullable)
512 }
513
514 fn nullable_i32(value: Option<i32>) -> Scalar {
515 match value {
516 Some(value) => Scalar::primitive::<i32>(value, Nullability::Nullable),
517 None => Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
518 }
519 }
520
521 fn ab_struct_dtype(nullability: Nullability) -> DType {
522 DType::Struct(
523 StructFields::new(
524 ["a", "b"].into(),
525 vec![
526 DType::Primitive(PType::I32, Nullability::NonNullable),
527 DType::Utf8(Nullability::NonNullable),
528 ],
529 ),
530 nullability,
531 )
532 }
533
534 #[rstest]
535 #[case(vec![0, 0], Some(true))]
537 #[case(vec![0], Some(true))]
538 #[case(vec![0, 5], Some(false))]
541 #[case(vec![5, 0], Some(false))]
542 #[case(vec![1, 2], Some(false))]
543 fn fixed_size_list_is_zero(#[case] values: Vec<i32>, #[case] expected: Option<bool>) {
544 let element_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
545 let children: Vec<Scalar> = values.into_iter().map(i32_scalar).collect();
546 let scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
547 assert_eq!(scalar.is_zero(), expected);
548 }
549
550 #[test]
551 fn null_fixed_size_list_is_zero_is_none() {
552 let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
553 let scalar = Scalar::null(DType::FixedSizeList(
554 element_dtype,
555 2,
556 Nullability::Nullable,
557 ));
558 assert_eq!(scalar.is_zero(), None);
559 }
560
561 #[test]
562 fn fixed_size_list_with_null_element_is_not_zero() {
563 let element_dtype = DType::Primitive(PType::I32, Nullability::Nullable);
566 let children = vec![nullable_i32(Some(0)), nullable_i32(None)];
567 let scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
568 assert_eq!(scalar.is_zero(), Some(false));
569 }
570
571 #[test]
572 fn struct_with_all_zero_fields_is_zero() {
573 let scalar = Scalar::struct_(
574 ab_struct_dtype(Nullability::NonNullable),
575 vec![i32_scalar(0), Scalar::utf8("", Nullability::NonNullable)],
576 );
577 assert_eq!(scalar.is_zero(), Some(true));
578 }
579
580 #[rstest]
581 #[case(5, "")]
584 #[case(0, "x")]
585 #[case(7, "y")]
586 fn struct_with_non_zero_field_is_not_zero(#[case] a: i32, #[case] b: &str) {
587 let scalar = Scalar::struct_(
588 ab_struct_dtype(Nullability::NonNullable),
589 vec![i32_scalar(a), Scalar::utf8(b, Nullability::NonNullable)],
590 );
591 assert_eq!(scalar.is_zero(), Some(false));
592 }
593
594 #[test]
595 fn null_struct_is_zero_is_none() {
596 let scalar = Scalar::null(ab_struct_dtype(Nullability::Nullable));
597 assert_eq!(scalar.is_zero(), None);
598 }
599
600 #[test]
601 fn struct_with_null_field_is_not_zero() {
602 let dtype = DType::Struct(
605 StructFields::new(
606 ["a", "b"].into(),
607 vec![
608 DType::Primitive(PType::I32, Nullability::Nullable),
609 DType::Primitive(PType::I32, Nullability::Nullable),
610 ],
611 ),
612 Nullability::NonNullable,
613 );
614 let scalar = Scalar::struct_(dtype, vec![nullable_i32(Some(0)), nullable_i32(None)]);
615 assert_eq!(scalar.is_zero(), Some(false));
616 }
617
618 #[test]
619 fn nested_struct_of_fixed_size_list_recurses() {
620 let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
623 let fsl_dtype =
624 DType::FixedSizeList(Arc::clone(&element_dtype), 2, Nullability::NonNullable);
625 let struct_dtype = DType::Struct(
626 StructFields::new(["fsl"].into(), vec![fsl_dtype]),
627 Nullability::NonNullable,
628 );
629
630 let all_zero = Scalar::struct_(
631 struct_dtype.clone(),
632 vec![Scalar::fixed_size_list(
633 Arc::clone(&element_dtype),
634 vec![i32_scalar(0), i32_scalar(0)],
635 Nullability::NonNullable,
636 )],
637 );
638 assert_eq!(all_zero.is_zero(), Some(true));
639
640 let with_non_zero = Scalar::struct_(
641 struct_dtype,
642 vec![Scalar::fixed_size_list(
643 element_dtype,
644 vec![i32_scalar(0), i32_scalar(9)],
645 Nullability::NonNullable,
646 )],
647 );
648 assert_eq!(with_non_zero.is_zero(), Some(false));
649 }
650}