Skip to main content

vortex_layout/layouts/row_idx/
expr.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Formatter;
5
6use vortex_array::ArrayRef;
7use vortex_array::dtype::DType;
8use vortex_array::dtype::Nullability;
9use vortex_array::dtype::PType;
10use vortex_array::expr::Expression;
11use vortex_array::scalar_fn::Arity;
12use vortex_array::scalar_fn::ChildName;
13use vortex_array::scalar_fn::EmptyOptions;
14use vortex_array::scalar_fn::ExecutionArgs;
15use vortex_array::scalar_fn::ScalarFnId;
16use vortex_array::scalar_fn::ScalarFnVTable;
17use vortex_array::scalar_fn::ScalarFnVTableExt;
18use vortex_error::VortexResult;
19use vortex_error::vortex_bail;
20
21#[derive(Clone)]
22pub struct RowIdx;
23
24impl ScalarFnVTable for RowIdx {
25    type Options = EmptyOptions;
26
27    fn id(&self) -> ScalarFnId {
28        ScalarFnId::from("vortex.row_idx")
29    }
30
31    fn arity(&self, _options: &Self::Options) -> Arity {
32        Arity::Exact(0)
33    }
34
35    fn child_name(&self, _instance: &Self::Options, _child_idx: usize) -> ChildName {
36        unreachable!()
37    }
38
39    fn fmt_sql(
40        &self,
41        _options: &Self::Options,
42        _expr: &Expression,
43        f: &mut Formatter<'_>,
44    ) -> std::fmt::Result {
45        write!(f, "#row_idx")
46    }
47
48    fn return_dtype(&self, _options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
49        Ok(DType::Primitive(PType::U64, Nullability::NonNullable))
50    }
51
52    fn execute(&self, _options: &Self::Options, _args: ExecutionArgs) -> VortexResult<ArrayRef> {
53        vortex_bail!(
54            "RowIdxExpr should not be executed directly, use it in the context of a Vortex scan and it will be substituted for a row index array"
55        );
56    }
57}
58
59pub fn row_idx() -> Expression {
60    RowIdx.new_expr(EmptyOptions, [])
61}