vortex_array/operator/
getitem.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::hash::{Hash, Hasher};
6use std::slice;
7use std::sync::Arc;
8
9use vortex_dtype::{DType, FieldName};
10use vortex_error::{VortexExpect, VortexResult};
11
12use crate::operator::{Operator, OperatorEq, OperatorHash, OperatorId, OperatorRef};
13
14/// An operator that extracts a field from a struct array.
15#[derive(Debug)]
16pub struct GetItemOperator {
17    // The struct-like child operator.
18    child: OperatorRef,
19    field: FieldName,
20    // The dtype of the extracted field.
21    dtype: DType,
22}
23
24impl OperatorHash for GetItemOperator {
25    fn operator_hash<H: Hasher>(&self, state: &mut H) {
26        self.child.operator_hash(state);
27        self.field.hash(state);
28        self.dtype.hash(state);
29    }
30}
31impl OperatorEq for GetItemOperator {
32    fn operator_eq(&self, other: &Self) -> bool {
33        self.child.operator_eq(&other.child)
34            && self.field == other.field
35            && self.dtype == other.dtype
36    }
37}
38
39impl GetItemOperator {
40    pub fn field_name(&self) -> &FieldName {
41        &self.field
42    }
43}
44
45impl Operator for GetItemOperator {
46    fn id(&self) -> OperatorId {
47        OperatorId::from("vortex.getitem")
48    }
49
50    fn as_any(&self) -> &dyn Any {
51        self
52    }
53
54    fn dtype(&self) -> &DType {
55        &self.dtype
56    }
57
58    fn len(&self) -> usize {
59        self.child.len()
60    }
61
62    fn children(&self) -> &[OperatorRef] {
63        slice::from_ref(&self.child)
64    }
65
66    fn with_children(self: Arc<Self>, children: Vec<OperatorRef>) -> VortexResult<OperatorRef> {
67        Ok(Arc::new(GetItemOperator {
68            child: children.into_iter().next().vortex_expect("missing child"),
69            field: self.field.clone(),
70            dtype: self.dtype.clone(),
71        }))
72    }
73}