vortex_layout/layouts/row_idx/
expr.rs1use 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}