lutra_arrow/
schema.rs

1use std::collections::HashSet;
2use std::ffi;
3use std::fs;
4use std::path;
5use std::sync::Arc;
6
7use lutra_bin::ident;
8use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
9use parquet::basic::Repetition;
10use parquet::file::metadata as pq_metadata;
11use parquet::schema::types as pq_types;
12
13#[derive(thiserror::Error, Debug)]
14pub enum Error {
15    #[error("{}: {}", .0.display(), .1)]
16    Io(std::path::PathBuf, std::io::Error),
17
18    #[error("invalid non-unicode path: {}", .0.display())]
19    NonUnicode(ffi::OsString),
20
21    #[error("parquet error: {0}")]
22    Parquet(#[from] parquet::errors::ParquetError),
23
24    #[error("arrow error: {0}")]
25    Arrow(#[from] arrow::error::ArrowError),
26
27    #[error("not implemented: {feature} at {location}")]
28    NotImplemented {
29        feature: &'static str,
30        location: String,
31    },
32}
33
34/// Scans the given base path for `.parquet` files and generates Lutra function
35/// definitions that use `std::fs::read_parquet()` to read them.
36///
37/// Does not traverse links.
38///
39/// Handles:
40/// - Recursive directory traversal
41/// - Parquet metadata reading
42/// - Type inference from schemas
43/// - Intelligent nullability detection
44pub fn pull_schema(base_path: &path::Path) -> Result<String, Error> {
45    let mut r = String::new();
46    let mut paths_to_search = vec![base_path.to_path_buf()];
47
48    while let Some(path) = paths_to_search.pop() {
49        let to_io_err = |e| Error::Io(path.clone(), e);
50
51        let md = path.metadata().map_err(to_io_err)?;
52        if md.is_dir() {
53            let read_dir = path.read_dir().map_err(to_io_err)?;
54            for entry in read_dir {
55                let entry = entry.map_err(to_io_err)?;
56                paths_to_search.push(entry.path());
57            }
58        } else if md.is_file() && path.extension().is_some_and(|e| e == "parquet") {
59            let relative = path.strip_prefix(base_path).unwrap();
60
61            r += &pull_schema_of_parquet_file(&path, relative)?;
62        }
63    }
64
65    Ok(r)
66}
67
68/// Read metadata of a parquet file and generates Lutra source code of a
69/// function that reads this file.
70fn pull_schema_of_parquet_file(
71    path: &path::Path,
72    path_relative: &path::Path,
73) -> Result<String, Error> {
74    // compose func name
75    let file_prefix = path.file_prefix().unwrap();
76    let func_name = to_str_or_err(file_prefix)?.replace(['/', ' '], "_");
77
78    // read parquet metadata
79    let file = fs::File::open(path).map_err(|e| Error::Io(path.to_path_buf(), e))?;
80    let mut metadata_reader = pq_metadata::ParquetMetaDataReader::new();
81    metadata_reader.try_parse(&file).unwrap();
82    let metadata = metadata_reader.finish().unwrap();
83
84    // Scan file to determine which optional fields actually have nulls
85    let file = fs::File::open(path).map_err(|e| Error::Io(path.to_path_buf(), e))?;
86    let nullable_columns = scan_for_nulls(file)?;
87
88    let ty = schema_to_lutra_ty(metadata.file_metadata().schema(), &nullable_columns)?;
89
90    // prepare path as str
91    let path = to_str_or_err(path_relative.as_os_str())?;
92
93    Ok(format!(
94        "\nfunc {func_name}(): {ty} -> std::fs::read_parquet(\"{path}\")\n"
95    ))
96}
97
98/// Scans the parquet file and returns the set of column names that contain null values.
99/// Note: This only checks top-level columns. Nested nullability is not checked.
100fn scan_for_nulls(file: fs::File) -> Result<HashSet<String>, Error> {
101    let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
102    let reader = builder.build()?;
103
104    let mut nullable_columns = HashSet::new();
105
106    for batch_result in reader {
107        let batch = batch_result?;
108
109        // Check each top-level column for null values
110        for (i, column) in batch.columns().iter().enumerate() {
111            if column.null_count() > 0 {
112                let field_name = batch.schema().field(i).name().clone();
113                nullable_columns.insert(field_name);
114            }
115        }
116    }
117
118    Ok(nullable_columns)
119}
120
121fn schema_to_lutra_ty(
122    schema: &pq_types::Type,
123    nullable_columns: &HashSet<String>,
124) -> Result<String, Error> {
125    let pq_types::Type::GroupType { fields, .. } = schema else {
126        return Err(Error::NotImplemented {
127            feature: "primitive type",
128            location: "top level".into(),
129        });
130    };
131
132    format_tuple_fields(fields, nullable_columns, "  ").map_err(|e| Error::NotImplemented {
133        feature: e.message,
134        location: e.location.unwrap_or_else(|| "top level".into()),
135    })
136}
137
138struct NotImplementedError {
139    message: &'static str,
140    location: Option<String>,
141}
142
143impl NotImplementedError {
144    fn new(message: &'static str) -> Self {
145        Self {
146            message,
147            location: None,
148        }
149    }
150
151    fn with_location(mut self, location: String) -> Self {
152        self.location = Some(location);
153        self
154    }
155}
156
157/// Format a list of parquet fields as a Lutra tuple type
158fn format_tuple_fields(
159    fields: &[Arc<pq_types::Type>],
160    nullable_columns: &HashSet<String>,
161    indent: &str,
162) -> Result<String, NotImplementedError> {
163    let mut ty_fields = String::new();
164
165    for f in fields {
166        let info = match f.as_ref() {
167            pq_types::Type::PrimitiveType { basic_info, .. } => basic_info,
168            pq_types::Type::GroupType { basic_info, .. } => basic_info,
169        };
170
171        // Get base type without considering repetition
172        let base_ty = field_base_type(f.as_ref(), indent)
173            .map_err(|e| e.with_location(format!("field {}", info.name())))?;
174
175        // Apply repetition and nullability
176        let field_ty = apply_repetition(
177            base_ty,
178            info.repetition(),
179            nullable_columns.contains(info.name()),
180        );
181
182        // Format field
183        ty_fields += indent;
184        ty_fields += &ident::display(info.name());
185        ty_fields += ": ";
186        ty_fields += &field_ty;
187        ty_fields += ",\n";
188    }
189
190    Ok(format!("[{{\n{ty_fields}}}]"))
191}
192
193/// Apply repetition and nullability wrapping to a base type
194fn apply_repetition(base_ty: String, repetition: Repetition, has_nulls: bool) -> String {
195    match repetition {
196        Repetition::OPTIONAL => {
197            if has_nulls {
198                format!("enum {{none, some: {base_ty}}}")
199            } else {
200                base_ty
201            }
202        }
203        Repetition::REQUIRED => base_ty,
204        Repetition::REPEATED => format!("[{base_ty}]"),
205    }
206}
207
208/// Get the base Lutra type for a field without considering repetition
209fn field_base_type(field: &pq_types::Type, indent: &str) -> Result<String, NotImplementedError> {
210    match field {
211        pq_types::Type::PrimitiveType {
212            basic_info,
213            physical_type,
214            ..
215        } => {
216            // Try converted type first, then fallback to physical type
217            let ty_field = converted_type_to_lt(basic_info.converted_type())?;
218            let ty_field = if let Some(t) = ty_field {
219                t.to_string()
220            } else {
221                physical_type_to_lt(*physical_type)?.to_string()
222            };
223            Ok(ty_field)
224        }
225        pq_types::Type::GroupType {
226            basic_info, fields, ..
227        } => {
228            // Check if this is a LIST type (special encoding)
229            if basic_info.converted_type() == parquet::basic::ConvertedType::LIST {
230                // LIST groups have a single field named "list" or "element"
231                if fields.len() == 1 {
232                    let indent = indent.to_string() + "  ";
233                    let element_ty = field_base_type(fields[0].as_ref(), &indent)?;
234                    return Ok(format!("[{element_ty}]"));
235                }
236            }
237
238            // Regular group type (nested struct/tuple)
239            // For nested structs, we don't check nullable_columns (would need path-based tracking)
240            // So we conservatively wrap all OPTIONAL nested fields
241            let empty_set = HashSet::new();
242            let indent = indent.to_string() + "  ";
243            format_tuple_fields(fields, &empty_set, &indent)
244        }
245    }
246}
247
248fn converted_type_to_lt(
249    ty: parquet::basic::ConvertedType,
250) -> Result<Option<&'static str>, NotImplementedError> {
251    use parquet::basic::ConvertedType;
252    Ok(Some(match ty {
253        ConvertedType::NONE => return Ok(None),
254
255        ConvertedType::UINT_8 => "Uint8",
256        ConvertedType::UINT_16 => "Uint16",
257        ConvertedType::UINT_32 => "Uint32",
258        ConvertedType::UINT_64 => "Uint64",
259        ConvertedType::INT_8 => "Int8",
260        ConvertedType::INT_16 => "Int16",
261        ConvertedType::INT_32 => "Int32",
262        ConvertedType::INT_64 => "Int64",
263        ConvertedType::UTF8 => "Text",
264
265        ConvertedType::MAP => return Err(NotImplementedError::new("MAP type")),
266        ConvertedType::MAP_KEY_VALUE => return Err(NotImplementedError::new("MAP_KEY_VALUE type")),
267        ConvertedType::LIST => return Ok(None), // Handled by caller
268        ConvertedType::ENUM => return Err(NotImplementedError::new("ENUM type")),
269        ConvertedType::DECIMAL => return Err(NotImplementedError::new("DECIMAL type")),
270
271        ConvertedType::DATE => "Date",
272        ConvertedType::TIME_MILLIS => return Err(NotImplementedError::new("TIME_MILLIS type")),
273        ConvertedType::TIME_MICROS => return Err(NotImplementedError::new("TIME_MICROS type")),
274        ConvertedType::TIMESTAMP_MILLIS => {
275            return Err(NotImplementedError::new("TIMESTAMP_MILLIS type"));
276        }
277        ConvertedType::TIMESTAMP_MICROS => "Timestamp",
278
279        ConvertedType::JSON => return Err(NotImplementedError::new("JSON type")),
280        ConvertedType::BSON => return Err(NotImplementedError::new("BSON type")),
281        ConvertedType::INTERVAL => return Err(NotImplementedError::new("INTERVAL type")),
282    }))
283}
284
285fn physical_type_to_lt(ty: parquet::basic::Type) -> Result<&'static str, NotImplementedError> {
286    use parquet::basic::Type;
287    Ok(match ty {
288        Type::BOOLEAN => "Bool",
289        Type::INT32 => "Int32",
290        Type::INT64 => "Int64",
291        Type::INT96 => return Err(NotImplementedError::new("INT96 type")),
292        Type::FLOAT => "Float32",
293        Type::DOUBLE => "Float64",
294        Type::BYTE_ARRAY => return Err(NotImplementedError::new("BYTE_ARRAY type")),
295        Type::FIXED_LEN_BYTE_ARRAY => {
296            return Err(NotImplementedError::new("FIXED_LEN_BYTE_ARRAY type"));
297        }
298    })
299}
300
301fn to_str_or_err(s: &ffi::OsStr) -> Result<&str, Error> {
302    s.to_str()
303        .ok_or_else(|| Error::NonUnicode(s.to_os_string()))
304}
lutra_arrow/schema.rs

lutra_arrow/
schema.rs
