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::expr::display::ExprDisplay;
12use vortex_array::scalar_fn::Arity;
13use vortex_array::scalar_fn::ChildName;
14use vortex_array::scalar_fn::EmptyOptions;
15use vortex_array::scalar_fn::ExecutionArgs;
16use vortex_array::scalar_fn::ScalarFnId;
17use vortex_array::scalar_fn::ScalarFnVTable;
18use vortex_array::scalar_fn::ScalarFnVTableExt;
19use vortex_error::VortexResult;
20use vortex_error::vortex_bail;
21use vortex_session::registry::CachedId;
22
23#[derive(Clone)]
24pub struct RowIdx;
25
26impl ScalarFnVTable for RowIdx {
27    type Options = EmptyOptions;
28
29    fn id(&self) -> ScalarFnId {
30        static ID: CachedId = CachedId::new("vortex.row_idx");
31        *ID
32    }
33
34    fn arity(&self, _options: &Self::Options) -> Arity {
35        Arity::Exact(0)
36    }
37
38    fn child_name(&self, _instance: &Self::Options, _child_idx: usize) -> ChildName {
39        unreachable!()
40    }
41
42    fn fmt_sql(
43        &self,
44        _options: &Self::Options,
45        _expr: &dyn ExprDisplay,
46        f: &mut Formatter<'_>,
47    ) -> std::fmt::Result {
48        write!(f, "#row_idx")
49    }
50
51    fn return_dtype(&self, _options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
52        Ok(DType::Primitive(PType::U64, Nullability::NonNullable))
53    }
54
55    fn execute(
56        &self,
57        _options: &Self::Options,
58        _args: &dyn ExecutionArgs,
59        _ctx: &mut vortex_array::ExecutionCtx,
60    ) -> VortexResult<ArrayRef> {
61        vortex_bail!(
62            "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"
63        );
64    }
65
66    fn is_strict(&self, _options: &Self::Options) -> bool {
67        true
68    }
69}
70
71pub fn row_idx() -> Expression {
72    RowIdx.new_expr(EmptyOptions, [])
73}