diff --git a/jitexpr/examples/basic.rs b/jitexpr/examples/basic.rs
index 8396c2a0c..6ecbb7865 100644
--- a/jitexpr/examples/basic.rs
+++ b/jitexpr/examples/basic.rs
@@ -1,24 +1,31 @@
 use std::collections::HashMap;
 use std::error::Error;
 
-use jitexpr::ast::{Function, InferredTypeSet, UntypedExpr, apply_types, infer_types};
+use jitexpr::ast::{Function, InferredTypeSet, UntypedExpr, infer_types};
 use jitexpr::compile::{CompiledFunction, compile};
 use jitexpr::types::VarType;
 
 fn main() -> Result<(), Box<dyn Error>> {
     // A simple expression that goes:
     // my_column + 1
+
     let untyped_expr = Function::Add.call_untyped_expr(vec![
         UntypedExpr::variable("my_col"),
         UntypedExpr::literal(1.0f64),
     ]);
 
-    let inferred_types = infer_types(&untyped_expr)?;
+    // Infer types does not return specific types, but instead a set of acceptable
+    // types for each variables.
+    let inferred_types: HashMap<&str, InferredTypeSet> = infer_types(&untyped_expr)?;
     assert_eq!(
         inferred_types.get("my_col").unwrap(),
         &InferredTypeSet::NUMERICAL
     );
 
+    // This is then up to us to decide the actual type for each variable.
+    // In tantivy, this means picking the first column with a type in inferred_types.
+    //
+    // If none match then we should use the VarType::None.
     let variable_types: HashMap<&str, VarType> =
         std::iter::once(("my_col", VarType::F64)).collect();
 
diff --git a/jitexpr/src/ast/mod.rs b/jitexpr/src/ast/mod.rs
index 0d202f12d..25ed8431f 100644
--- a/jitexpr/src/ast/mod.rs
+++ b/jitexpr/src/ast/mod.rs
@@ -1,15 +1,13 @@
-mod apply_types;
 mod infer_types;
 mod literal;
-mod typed_expr;
 mod untyped_expr;
 
-pub use apply_types::apply_types;
 pub use infer_types::{InferredTypeSet, infer_types};
 pub use literal::Literal;
-pub use typed_expr::{TypedExpr, TypedExprAst, TypedVariable};
 pub use untyped_expr::UntypedExpr;
 
+use crate::compile::{TypedExpr, TypedExprAst};
+
 /// A function supported by the first expression-language milestone.
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 pub enum Function {
diff --git a/jitexpr/src/ast/apply_types.rs b/jitexpr/src/compile/apply_types.rs
similarity index 96%
rename from jitexpr/src/ast/apply_types.rs
rename to jitexpr/src/compile/apply_types.rs
index 77e6b0745..352d4b492 100644
--- a/jitexpr/src/ast/apply_types.rs
+++ b/jitexpr/src/compile/apply_types.rs
@@ -1,8 +1,8 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::ast::typed_expr::TypedVariable;
-use crate::ast::{Function, Literal, TypedExpr, TypedExprAst, UntypedExpr};
+use super::typed_expr::{TypedExpr, TypedExprAst, TypedVariable};
+use crate::ast::{Function, Literal, UntypedExpr};
 use crate::types::VarType;
 
 /// If a variable is missing from variable_types, it will be treated as if its value is None.
@@ -25,18 +25,22 @@ fn apply_types_aux(
             ast: TypedExprAst::Literal(literal.clone()),
         },
         UntypedExpr::Variable(variable_name) => {
-            if let Some(variable_type) = variable_types.get(variable_name.as_ref()).copied() {
-                TypedExpr {
-                    return_type: variable_type,
-                    ast: TypedExprAst::variable(variable_name, variable_type),
-                }
-            } else {
+            let variable_type: VarType = variable_types
+                .get(variable_name.as_ref())
+                .copied()
                 // a missing column is treated as if it was there with a constant
                 // None value.
+                .unwrap_or(VarType::None);
+            if variable_type == VarType::None {
                 TypedExpr {
                     return_type: VarType::None,
                     ast: TypedExprAst::Literal(Literal::None),
                 }
+            } else {
+                TypedExpr {
+                    return_type: variable_type,
+                    ast: TypedExprAst::variable(variable_name, variable_type),
+                }
             }
         }
         UntypedExpr::Call { function, args } => match function {
diff --git a/jitexpr/src/compile/mod.rs b/jitexpr/src/compile/mod.rs
index 64c3c6ac1..7bf99d7db 100644
--- a/jitexpr/src/compile/mod.rs
+++ b/jitexpr/src/compile/mod.rs
@@ -1,14 +1,17 @@
+mod apply_types;
+mod typed_expr;
+
 use std::collections::HashMap;
 use std::mem::{self, size_of};
 
+pub use apply_types::apply_types;
 use cranelift::codegen::ir::{MemFlagsData, UserFuncName};
 use cranelift::prelude::*;
 use cranelift_jit::{JITBuilder, JITModule};
 use cranelift_module::{Module, default_libcall_names};
+pub use typed_expr::{TypedExpr, TypedExprAst, TypedVariable};
 
-use crate::ast::{
-    Function, Literal, TypedExpr, TypedExprAst, TypedVariable, UntypedExpr, apply_types,
-};
+use crate::ast::{Function, Literal, UntypedExpr};
 use crate::types::{StringRef, VarType, VariableValue};
 
 /// An expression compiled to native machine code.
diff --git a/jitexpr/src/ast/typed_expr.rs b/jitexpr/src/compile/typed_expr.rs
similarity index 97%
rename from jitexpr/src/ast/typed_expr.rs
rename to jitexpr/src/compile/typed_expr.rs
index d83d3fa1e..67fcebf5d 100644
--- a/jitexpr/src/ast/typed_expr.rs
+++ b/jitexpr/src/compile/typed_expr.rs
@@ -6,8 +6,8 @@ use crate::types::VarType;
 #[derive(Clone, PartialEq)]
 pub struct TypedVariable {
     pub(super) variable_name: Arc<str>,
-    pub(crate) r#type: VarType,
-    pub(crate) variable_id: usize, //< offset in the input array.
+    pub(super) r#type: VarType,
+    pub(super) variable_id: usize, //< offset in the input array.
 }
 
 #[derive(Clone, PartialEq)]