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