lutra_arrow/
interface.rs

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

lutra_arrow/
interface.rs
