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 {
152 let value = ScalarValue::zero_value(dtype);
153
154 unsafe { Self::new_unchecked(dtype.clone(), Some(value)) }
156 }
157
158 pub fn eq_ignore_nullability(&self, other: &Self) -> bool {
162 self.dtype.eq_ignore_nullability(&other.dtype) && self.value == other.value
163 }
164
165 pub fn into_parts(self) -> (DType, Option<ScalarValue>) {
167 (self.dtype, self.value)
168 }
169
170 pub fn dtype(&self) -> &DType {
172 &self.dtype
173 }
174
175 pub fn value(&self) -> Option<&ScalarValue> {
177 self.value.as_ref()
178 }
179
180 pub fn into_value(self) -> Option<ScalarValue> {
183 self.value
184 }
185
186 pub fn is_valid(&self) -> bool {
188 self.value.is_some()
189 }
190
191 pub fn is_null(&self) -> bool {
193 self.value.is_none()
194 }
195
196 pub fn is_zero(&self) -> Option<bool> {
205 let value = self.value()?;
206
207 let is_zero = match self.dtype() {
208 DType::Null => vortex_panic!("non-null value somehow had `DType::Null`"),
209 DType::Bool(_) => !value.as_bool(),
210 DType::Primitive(..) => value.as_primitive().is_zero(),
211 DType::Decimal(..) => value.as_decimal().is_zero(),
212 DType::Utf8(_) => value.as_utf8().is_empty(),
213 DType::Binary(_) => value.as_binary().is_empty(),
214 DType::List(..) => value.as_list().is_empty(),
215 DType::FixedSizeList(_, list_size, _) => {
218 let list = self.as_list();
219 list.len() == *list_size as usize
220 && (0..list.len())
221 .all(|i| list.element(i).is_some_and(|e| e.is_zero() == Some(true)))
222 }
223 DType::Struct(..) => self
225 .as_struct()
226 .fields_iter()
227 .is_some_and(|mut fields| fields.all(|f| f.is_zero() == Some(true))),
228 DType::Union(..) => {
231 let union = self.as_union();
232 if union.child_index() != Some(0) {
233 false
234 } else {
235 union.child().and_then(|child| child.is_zero())?
236 }
237 }
238 DType::Variant(_) => self.as_variant().is_zero()?,
239 DType::Extension(_) => self.as_extension().to_storage_scalar().is_zero()?,
240 };
241
242 Some(is_zero)
243 }
244
245 pub fn primitive_reinterpret_cast(&self, ptype: PType) -> VortexResult<Self> {
251 let primitive = self.as_primitive();
252 if primitive.ptype() == ptype {
253 return Ok(self.clone());
254 }
255
256 vortex_ensure_eq!(
257 primitive.ptype().byte_width(),
258 ptype.byte_width(),
259 "can't reinterpret cast between integers of two different widths"
260 );
261
262 Scalar::try_new(
263 DType::Primitive(ptype, self.dtype().nullability()),
264 primitive
265 .pvalue()
266 .map(|p| p.reinterpret_cast(ptype))
267 .map(ScalarValue::Primitive),
268 )
269 }
270
271 pub fn approx_nbytes(&self) -> usize {
276 use crate::dtype::NativeDecimalType;
277 use crate::dtype::i256;
278
279 match self.dtype() {
280 DType::Null => 0,
281 DType::Bool(_) => 1,
282 DType::Primitive(ptype, _) => ptype.byte_width(),
283 DType::Decimal(dt, _) => {
284 if dt.precision() <= i128::MAX_PRECISION {
285 size_of::<i128>()
286 } else {
287 size_of::<i256>()
288 }
289 }
290 DType::Utf8(_) => self
291 .value()
292 .map_or_else(|| 0, |value| value.as_utf8().len()),
293 DType::Binary(_) => self
294 .value()
295 .map_or_else(|| 0, |value| value.as_binary().len()),
296 DType::List(..) | DType::FixedSizeList(..) => self
297 .as_list()
298 .elements()
299 .map(|fields| fields.into_iter().map(|f| f.approx_nbytes()).sum::<usize>())
300 .unwrap_or_default(),
301 DType::Struct(..) => self
302 .as_struct()
303 .fields_iter()
304 .map(|fields| fields.into_iter().map(|f| f.approx_nbytes()).sum::<usize>())
305 .unwrap_or_default(),
306 DType::Union(..) => self
307 .as_union()
308 .child()
309 .map_or(0, |value| 1 + value.approx_nbytes()),
310 DType::Variant(_) => self.as_variant().value().map_or(0, Scalar::approx_nbytes),
311 DType::Extension(_) => self.as_extension().to_storage_scalar().approx_nbytes(),
312 }
313 }
314}
315
316impl Hash for Scalar {
320 fn hash<H: Hasher>(&self, state: &mut H) {
321 self.dtype.hash_ignore_nullability(state);
322 self.value.hash(state);
323 }
324}
325
326impl PartialEq for Scalar {
333 fn eq(&self, other: &Self) -> bool {
334 self.dtype.eq_ignore_nullability(&other.dtype) && self.value == other.value
335 }
336}
337
338impl PartialOrd for Scalar {
339 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
375 if !self.dtype().eq_ignore_nullability(other.dtype()) {
376 return None;
377 }
378
379 partial_cmp_scalar_values(self.dtype(), self.value(), other.value())
380 }
381}
382
383fn partial_cmp_scalar_values(
385 dtype: &DType,
386 lhs: Option<&ScalarValue>,
387 rhs: Option<&ScalarValue>,
388) -> Option<Ordering> {
389 match (lhs, rhs) {
390 (None, None) => Some(Ordering::Equal),
391 (None, Some(_)) => Some(Ordering::Less),
392 (Some(_), None) => Some(Ordering::Greater),
393 (Some(lhs), Some(rhs)) => partial_cmp_non_null_scalar_values(dtype, lhs, rhs),
394 }
395}
396
397fn partial_cmp_non_null_scalar_values(
399 dtype: &DType,
400 lhs: &ScalarValue,
401 rhs: &ScalarValue,
402) -> Option<Ordering> {
403 match (lhs, rhs) {
406 (ScalarValue::Bool(lhs), ScalarValue::Bool(rhs)) => lhs.partial_cmp(rhs),
407 (ScalarValue::Primitive(lhs), ScalarValue::Primitive(rhs)) => lhs.partial_cmp(rhs),
408 (ScalarValue::Decimal(lhs), ScalarValue::Decimal(rhs)) => lhs.partial_cmp(rhs),
409 (ScalarValue::Utf8(lhs), ScalarValue::Utf8(rhs)) => lhs.partial_cmp(rhs),
410 (ScalarValue::Binary(lhs), ScalarValue::Binary(rhs)) => lhs.partial_cmp(rhs),
411 (ScalarValue::Tuple(lhs), ScalarValue::Tuple(rhs)) => {
414 partial_cmp_tuple_values(dtype, lhs, rhs)
415 }
416 (ScalarValue::Union(lhs), ScalarValue::Union(rhs)) => {
417 if lhs.type_id() != rhs.type_id() {
418 return None;
419 }
420
421 let DType::Union(variants, _) = dtype else {
422 return None;
423 };
424 let child_index = variants.tag_to_child_index(lhs.type_id())?;
425 let child_dtype = variants.variant_by_index(child_index)?;
426
427 partial_cmp_scalar_values(&child_dtype, lhs.child_value(), rhs.child_value())
428 }
429 (ScalarValue::Variant(_), ScalarValue::Variant(_)) => None,
432 _ => None,
433 }
434}
435
436fn partial_cmp_tuple_values(
438 dtype: &DType,
439 lhs: &[Option<ScalarValue>],
440 rhs: &[Option<ScalarValue>],
441) -> Option<Ordering> {
442 match dtype {
443 DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => {
444 partial_cmp_list_values(element_dtype, lhs, rhs)
445 }
446 DType::Struct(fields, _) => partial_cmp_struct_values(fields, lhs, rhs),
447 DType::Extension(ext_dtype) => {
448 partial_cmp_tuple_values(ext_dtype.storage_dtype(), lhs, rhs)
449 }
450 _ => None,
451 }
452}
453
454fn partial_cmp_list_values(
456 element_dtype: &DType,
457 lhs: &[Option<ScalarValue>],
458 rhs: &[Option<ScalarValue>],
459) -> Option<Ordering> {
460 for (lhs, rhs) in lhs.iter().zip(rhs.iter()) {
461 match partial_cmp_scalar_values(element_dtype, lhs.as_ref(), rhs.as_ref())? {
462 Ordering::Equal => continue,
463 ordering => return Some(ordering),
464 }
465 }
466
467 Some(lhs.len().cmp(&rhs.len()))
468}
469
470fn partial_cmp_struct_values(
472 fields: &StructFields,
473 lhs: &[Option<ScalarValue>],
474 rhs: &[Option<ScalarValue>],
475) -> Option<Ordering> {
476 if lhs.len() != fields.nfields() || rhs.len() != fields.nfields() {
477 return None;
478 }
479
480 for ((field_dtype, lhs), rhs) in fields.fields().zip(lhs.iter()).zip(rhs.iter()) {
481 match partial_cmp_scalar_values(&field_dtype, lhs.as_ref(), rhs.as_ref())? {
482 Ordering::Equal => continue,
483 ordering => return Some(ordering),
484 }
485 }
486
487 Some(Ordering::Equal)
488}
489
490#[cfg(test)]
491mod tests {
492 use std::sync::Arc;
493
494 use rstest::rstest;
495
496 use crate::dtype::DType;
497 use crate::dtype::Nullability;
498 use crate::dtype::PType;
499 use crate::dtype::StructFields;
500 use crate::scalar::Scalar;
501
502 fn i32_scalar(value: i32) -> Scalar {
503 Scalar::primitive::<i32>(value, Nullability::NonNullable)
504 }
505
506 fn nullable_i32(value: Option<i32>) -> Scalar {
507 match value {
508 Some(value) => Scalar::primitive::<i32>(value, Nullability::Nullable),
509 None => Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
510 }
511 }
512
513 fn ab_struct_dtype(nullability: Nullability) -> DType {
514 DType::Struct(
515 StructFields::new(
516 ["a", "b"].into(),
517 vec![
518 DType::Primitive(PType::I32, Nullability::NonNullable),
519 DType::Utf8(Nullability::NonNullable),
520 ],
521 ),
522 nullability,
523 )
524 }
525
526 #[rstest]
527 #[case(vec![0, 0], Some(true))]
529 #[case(vec![0], Some(true))]
530 #[case(vec![0, 5], Some(false))]
533 #[case(vec![5, 0], Some(false))]
534 #[case(vec![1, 2], Some(false))]
535 fn fixed_size_list_is_zero(#[case] values: Vec<i32>, #[case] expected: Option<bool>) {
536 let element_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
537 let children: Vec<Scalar> = values.into_iter().map(i32_scalar).collect();
538 let scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
539 assert_eq!(scalar.is_zero(), expected);
540 }
541
542 #[test]
543 fn null_fixed_size_list_is_zero_is_none() {
544 let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
545 let scalar = Scalar::null(DType::FixedSizeList(
546 element_dtype,
547 2,
548 Nullability::Nullable,
549 ));
550 assert_eq!(scalar.is_zero(), None);
551 }
552
553 #[test]
554 fn fixed_size_list_with_null_element_is_not_zero() {
555 let element_dtype = DType::Primitive(PType::I32, Nullability::Nullable);
558 let children = vec![nullable_i32(Some(0)), nullable_i32(None)];
559 let scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
560 assert_eq!(scalar.is_zero(), Some(false));
561 }
562
563 #[test]
564 fn struct_with_all_zero_fields_is_zero() {
565 let scalar = Scalar::struct_(
566 ab_struct_dtype(Nullability::NonNullable),
567 vec![i32_scalar(0), Scalar::utf8("", Nullability::NonNullable)],
568 );
569 assert_eq!(scalar.is_zero(), Some(true));
570 }
571
572 #[rstest]
573 #[case(5, "")]
576 #[case(0, "x")]
577 #[case(7, "y")]
578 fn struct_with_non_zero_field_is_not_zero(#[case] a: i32, #[case] b: &str) {
579 let scalar = Scalar::struct_(
580 ab_struct_dtype(Nullability::NonNullable),
581 vec![i32_scalar(a), Scalar::utf8(b, Nullability::NonNullable)],
582 );
583 assert_eq!(scalar.is_zero(), Some(false));
584 }
585
586 #[test]
587 fn null_struct_is_zero_is_none() {
588 let scalar = Scalar::null(ab_struct_dtype(Nullability::Nullable));
589 assert_eq!(scalar.is_zero(), None);
590 }
591
592 #[test]
593 fn struct_with_null_field_is_not_zero() {
594 let dtype = DType::Struct(
597 StructFields::new(
598 ["a", "b"].into(),
599 vec![
600 DType::Primitive(PType::I32, Nullability::Nullable),
601 DType::Primitive(PType::I32, Nullability::Nullable),
602 ],
603 ),
604 Nullability::NonNullable,
605 );
606 let scalar = Scalar::struct_(dtype, vec![nullable_i32(Some(0)), nullable_i32(None)]);
607 assert_eq!(scalar.is_zero(), Some(false));
608 }
609
610 #[test]
611 fn nested_struct_of_fixed_size_list_recurses() {
612 let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable));
615 let fsl_dtype =
616 DType::FixedSizeList(Arc::clone(&element_dtype), 2, Nullability::NonNullable);
617 let struct_dtype = DType::Struct(
618 StructFields::new(["fsl"].into(), vec![fsl_dtype]),
619 Nullability::NonNullable,
620 );
621
622 let all_zero = Scalar::struct_(
623 struct_dtype.clone(),
624 vec![Scalar::fixed_size_list(
625 Arc::clone(&element_dtype),
626 vec![i32_scalar(0), i32_scalar(0)],
627 Nullability::NonNullable,
628 )],
629 );
630 assert_eq!(all_zero.is_zero(), Some(true));
631
632 let with_non_zero = Scalar::struct_(
633 struct_dtype,
634 vec![Scalar::fixed_size_list(
635 element_dtype,
636 vec![i32_scalar(0), i32_scalar(9)],
637 Nullability::NonNullable,
638 )],
639 );
640 assert_eq!(with_non_zero.is_zero(), Some(false));
641 }
642}