Skip to main content

lutra_arrow/
context.rs

1use lutra_bin::ir;
2use std::collections::HashMap;
3
4/// Context for resolving type identifiers
5pub struct Context<'a> {
6    types: HashMap<&'a ir::Path, &'a ir::Ty>,
7}
8
9impl<'a> Context<'a> {
10    pub fn new(ty_defs: &'a [ir::TyDef]) -> Self {
11        Context {
12            types: ty_defs.iter().map(|def| (&def.name, &def.ty)).collect(),
13        }
14    }
15
16    /// Resolve type identifiers to their materialized types
17    pub fn get_ty_mat(&self, ty: &'a ir::Ty) -> Result<&'a ir::Ty, Error> {
18        let mut ty = ty;
19        while let ir::TyKind::Ident(path) = &ty.kind {
20            ty = self.types.get(path).ok_or(Error::BadType)?;
21        }
22        Ok(ty)
23    }
24
25    /// Checks if enum is option pattern: enum {none, some: T} where T is primitive or array.
26    /// These can be represented as nullable values (single column).
27    pub fn is_option(&self, variants: &[ir::TyEnumVariant]) -> Result<bool, Error> {
28        if variants.len() != 2 || !variants[0].ty.is_unit() {
29            return Ok(false);
30        }
31        let some_ty = self.get_ty_mat(&variants[1].ty)?;
32        Ok(some_ty.kind.is_primitive() || some_ty.kind.is_array())
33    }
34}
35
36/// Common error type for Arrow conversion operations
37#[derive(Debug, thiserror::Error)]
38pub enum Error {
39    #[error("type mismatch: Arrow has {arrow_ty}, but Lutra expects {lutra_ty}")]
40    TypeMismatch { arrow_ty: String, lutra_ty: String },
41
42    #[error("expected single row, got {got}")]
43    ExpectedSingleRow { got: usize },
44
45    #[error("expected {expected} columns, got {got}")]
46    ColumnCountMismatch { expected: usize, got: usize },
47
48    #[error("unexpected null value for non-option type")]
49    UnexpectedNull,
50
51    #[error("provided type is invalid")]
52    BadType,
53
54    #[error("unsupported type for Arrow conversion: {0}")]
55    UnsupportedType(String),
56
57    #[error("Arrow error: {0}")]
58    Arrow(#[from] arrow::error::ArrowError),
59}