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    /// Single-step ident resolution (one hop only)
17    pub fn get_ty(&self, ident: &'a ir::Path) -> Result<&'a ir::Ty, Error> {
18        self.types.get(ident).copied().ok_or(Error::BadType)
19    }
20
21    /// Resolve type identifiers to their materialized types
22    pub fn get_ty_mat(&self, ty: &'a ir::Ty) -> Result<&'a ir::Ty, Error> {
23        let mut ty = ty;
24        while let ir::TyKind::Ident(path) = &ty.kind {
25            if ir::TyStd::try_new(path).is_some() {
26                return Ok(ty);
27            }
28            ty = self.types.get(path).ok_or(Error::BadType)?;
29        }
30        Ok(ty)
31    }
32
33    /// Checks if enum is option pattern: enum {none, some: T} where T is primitive or array.
34    /// These can be represented as nullable values (single column).
35    pub fn is_option(&self, variants: &[ir::TyEnumVariant]) -> Result<bool, Error> {
36        if variants.len() != 2 || !variants[0].ty.is_unit() {
37            return Ok(false);
38        }
39        let some_ty = self.get_ty_mat(&variants[1].ty)?;
40        Ok(some_ty.kind.is_primitive() || some_ty.kind.is_ident() || some_ty.kind.is_array())
41    }
42}
43
44/// Common error type for Arrow conversion operations
45#[derive(Debug, thiserror::Error)]
46pub enum Error {
47    #[error("type mismatch: Arrow has {arrow_ty}, but Lutra expects {lutra_ty}")]
48    TypeMismatch { arrow_ty: String, lutra_ty: String },
49
50    #[error("expected single row, got {got}")]
51    ExpectedSingleRow { got: usize },
52
53    #[error("expected {expected} columns, got {got}")]
54    ColumnCountMismatch { expected: usize, got: usize },
55
56    #[error("unexpected null value for non-option type")]
57    UnexpectedNull,
58
59    #[error("provided type is invalid")]
60    BadType,
61
62    #[error("unsupported type for Arrow conversion: {0}")]
63    UnsupportedType(String),
64
65    #[error("Arrow error: {0}")]
66    Arrow(#[from] arrow::error::ArrowError),
67}