1use std::{fs, path::Path, sync::Arc};
2
3use anyhow::{Context, Result};
4use arrow_array::{ArrayRef, Float64Array, Int64Array, RecordBatch, StringArray};
5use arrow_ipc::writer::FileWriter;
6use arrow_schema::{DataType, Field, Schema};
7use rust_xlsxwriter::{Format, Workbook};
8
9use crate::{QueryResult, QueryValue};
10
11pub fn write_xlsx(result: &QueryResult, output_path: &Path) -> Result<()> {
12 const MIN_COLUMN_WIDTH: f64 = 8.0;
13 const MAX_COLUMN_WIDTH: f64 = 60.0;
14 const COLUMN_PADDING: f64 = 2.0;
15
16 let mut workbook = Workbook::new();
17 let worksheet = workbook.add_worksheet();
18 let header_format = Format::new().set_bold();
19
20 let mut max_widths: Vec<usize> = result
21 .columns
22 .iter()
23 .map(|header| header.chars().count())
24 .collect();
25
26 for (column, header) in result.columns.iter().enumerate() {
27 worksheet.write_string_with_format(0, column as u16, header, &header_format)?;
28 }
29
30 for (row_index, row) in result.rows.iter().enumerate() {
31 for (column_index, value) in row.iter().enumerate() {
32 let excel_row = (row_index + 1) as u32;
33 let excel_column = column_index as u16;
34
35 if let Some(current_max) = max_widths.get_mut(column_index) {
36 *current_max = (*current_max).max(query_value_display_width(value));
37 }
38
39 match value {
40 QueryValue::Null => {}
41 QueryValue::Integer(number) => {
42 worksheet.write_number(excel_row, excel_column, *number as f64)?;
43 }
44 QueryValue::Real(number) => {
45 worksheet.write_number(excel_row, excel_column, *number)?;
46 }
47 QueryValue::Text(text) => {
48 worksheet.write_string(excel_row, excel_column, text)?;
49 }
50 }
51 }
52 }
53
54 for (column_index, max_width) in max_widths.iter().enumerate() {
55 let width = ((*max_width as f64) + COLUMN_PADDING).clamp(MIN_COLUMN_WIDTH, MAX_COLUMN_WIDTH);
56 worksheet.set_column_width(column_index as u16, width)?;
57 }
58
59 if !result.columns.is_empty() {
60 let last_column = (result.columns.len() - 1) as u16;
61 let last_row = result.rows.len() as u32;
62 worksheet.autofilter(0, 0, last_row, last_column)?;
63 }
64
65 workbook
66 .save(output_path)
67 .with_context(|| format!("failed to write {}", output_path.display()))
68}
69
70fn query_value_display_width(value: &QueryValue) -> usize {
71 match value {
72 QueryValue::Null => 0,
73 QueryValue::Integer(number) => number.to_string().chars().count(),
74 QueryValue::Real(number) => number.to_string().chars().count(),
75 QueryValue::Text(text) => text.chars().count(),
76 }
77}
78
79pub fn write_parquet(result: &QueryResult, output_path: &Path) -> Result<()> {
80 use parquet::basic::{ConvertedType, Repetition, Type as PhysicalType};
81 use parquet::column::writer::ColumnWriter;
82 use parquet::data_type::ByteArray;
83 use parquet::file::properties::WriterProperties;
84 use parquet::file::writer::SerializedFileWriter;
85 use parquet::schema::types::Type as ParquetType;
86
87 #[derive(Clone, Copy)]
88 enum ColKind {
89 Int64,
90 Double,
91 Bytes,
92 }
93
94 let col_kinds: Vec<ColKind> = (0..result.columns.len())
96 .map(|ci| {
97 let mut has_real = false;
98 let mut has_text = false;
99 for row in &result.rows {
100 match row.get(ci).unwrap_or(&QueryValue::Null) {
101 QueryValue::Real(_) => has_real = true,
102 QueryValue::Text(_) => has_text = true,
103 _ => {}
104 }
105 }
106 if has_text {
107 ColKind::Bytes
108 } else if has_real {
109 ColKind::Double
110 } else {
111 ColKind::Int64
112 }
113 })
114 .collect();
115
116 let fields: Vec<Arc<ParquetType>> = result
118 .columns
119 .iter()
120 .zip(&col_kinds)
121 .map(|(name, kind)| {
122 let physical = match kind {
123 ColKind::Int64 => PhysicalType::INT64,
124 ColKind::Double => PhysicalType::DOUBLE,
125 ColKind::Bytes => PhysicalType::BYTE_ARRAY,
126 };
127 let mut builder = ParquetType::primitive_type_builder(name, physical)
128 .with_repetition(Repetition::OPTIONAL);
129 if matches!(kind, ColKind::Bytes) {
130 builder = builder.with_converted_type(ConvertedType::UTF8);
131 }
132 Arc::new(
133 builder
134 .build()
135 .with_context(|| format!("failed to build Parquet field '{name}'"))
136 .expect("valid field"),
137 )
138 })
139 .collect();
140
141 let schema = Arc::new(
142 ParquetType::group_type_builder("schema")
143 .with_fields(fields)
144 .build()
145 .context("failed to build Parquet schema")?,
146 );
147
148 let props = Arc::new(WriterProperties::builder().build());
149 let file = fs::File::create(output_path)
150 .with_context(|| format!("failed to create {}", output_path.display()))?;
151 let mut file_writer = SerializedFileWriter::new(file, schema, props)
152 .context("failed to initialize Parquet writer")?;
153
154 let mut rg = file_writer
155 .next_row_group()
156 .context("failed to start Parquet row group")?;
157
158 for (col_idx, kind) in col_kinds.iter().enumerate() {
159 let def_levels: Vec<i16> = result
160 .rows
161 .iter()
162 .map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
163 QueryValue::Null => 0,
164 _ => 1,
165 })
166 .collect();
167
168 let Some(mut col_writer) = rg
169 .next_column()
170 .context("failed to open Parquet column writer")?
171 else {
172 break;
173 };
174
175 match (kind, col_writer.untyped()) {
176 (ColKind::Int64, ColumnWriter::Int64ColumnWriter(w)) => {
177 let values: Vec<i64> = result
178 .rows
179 .iter()
180 .filter_map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
181 QueryValue::Integer(v) => Some(*v),
182 _ => None,
183 })
184 .collect();
185 w.write_batch(&values, Some(&def_levels), None)?;
186 }
187 (ColKind::Double, ColumnWriter::DoubleColumnWriter(w)) => {
188 let values: Vec<f64> = result
189 .rows
190 .iter()
191 .filter_map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
192 QueryValue::Real(v) => Some(*v),
193 QueryValue::Integer(v) => Some(*v as f64),
194 _ => None,
195 })
196 .collect();
197 w.write_batch(&values, Some(&def_levels), None)?;
198 }
199 (ColKind::Bytes, ColumnWriter::ByteArrayColumnWriter(w)) => {
200 let values: Vec<ByteArray> = result
201 .rows
202 .iter()
203 .filter_map(|row| match row.get(col_idx).unwrap_or(&QueryValue::Null) {
204 QueryValue::Text(s) => Some(ByteArray::from(s.as_bytes().to_vec())),
205 QueryValue::Integer(v) => Some(ByteArray::from(v.to_string().into_bytes())),
206 QueryValue::Real(v) => Some(ByteArray::from(v.to_string().into_bytes())),
207 _ => None,
208 })
209 .collect();
210 w.write_batch(&values, Some(&def_levels), None)?;
211 }
212 _ => unreachable!("column kind and writer type must agree"),
213 }
214
215 col_writer.close()?;
216 }
217
218 rg.close()?;
219 file_writer.close()?;
220
221 Ok(())
222}
223
224pub fn write_feather(result: &QueryResult, output_path: &Path) -> Result<()> {
225 #[derive(Clone, Copy)]
226 enum ColKind {
227 Int64,
228 Double,
229 Utf8,
230 }
231
232 let col_kinds: Vec<ColKind> = (0..result.columns.len())
233 .map(|column_index| {
234 let mut has_real = false;
235 let mut has_text = false;
236 for row in &result.rows {
237 match row.get(column_index).unwrap_or(&QueryValue::Null) {
238 QueryValue::Real(_) => has_real = true,
239 QueryValue::Text(_) => has_text = true,
240 _ => {}
241 }
242 }
243 if has_text {
244 ColKind::Utf8
245 } else if has_real {
246 ColKind::Double
247 } else {
248 ColKind::Int64
249 }
250 })
251 .collect();
252
253 let schema = Arc::new(Schema::new(
254 result
255 .columns
256 .iter()
257 .zip(&col_kinds)
258 .map(|(name, kind)| {
259 let data_type = match kind {
260 ColKind::Int64 => DataType::Int64,
261 ColKind::Double => DataType::Float64,
262 ColKind::Utf8 => DataType::Utf8,
263 };
264 Field::new(name, data_type, true)
265 })
266 .collect::<Vec<_>>(),
267 ));
268
269 let arrays = col_kinds
270 .iter()
271 .enumerate()
272 .map(|(column_index, kind)| match kind {
273 ColKind::Int64 => {
274 let values = result
275 .rows
276 .iter()
277 .map(
278 |row| match row.get(column_index).unwrap_or(&QueryValue::Null) {
279 QueryValue::Integer(value) => Some(*value),
280 QueryValue::Null => None,
281 QueryValue::Real(_) | QueryValue::Text(_) => None,
282 },
283 )
284 .collect::<Vec<_>>();
285 Arc::new(Int64Array::from(values)) as ArrayRef
286 }
287 ColKind::Double => {
288 let values = result
289 .rows
290 .iter()
291 .map(
292 |row| match row.get(column_index).unwrap_or(&QueryValue::Null) {
293 QueryValue::Integer(value) => Some(*value as f64),
294 QueryValue::Real(value) => Some(*value),
295 QueryValue::Null | QueryValue::Text(_) => None,
296 },
297 )
298 .collect::<Vec<_>>();
299 Arc::new(Float64Array::from(values)) as ArrayRef
300 }
301 ColKind::Utf8 => {
302 let values = result
303 .rows
304 .iter()
305 .map(
306 |row| match row.get(column_index).unwrap_or(&QueryValue::Null) {
307 QueryValue::Integer(value) => Some(value.to_string()),
308 QueryValue::Real(value) => Some(value.to_string()),
309 QueryValue::Text(value) => Some(value.clone()),
310 QueryValue::Null => None,
311 },
312 )
313 .collect::<Vec<_>>();
314 Arc::new(StringArray::from(values)) as ArrayRef
315 }
316 })
317 .collect::<Vec<_>>();
318
319 let batch = RecordBatch::try_new(schema.clone(), arrays)
320 .context("failed to build Feather record batch")?;
321 let file = fs::File::create(output_path)
322 .with_context(|| format!("failed to create {}", output_path.display()))?;
323 let mut writer =
324 FileWriter::try_new(file, &schema).context("failed to initialize Feather writer")?;
325 writer
326 .write(&batch)
327 .context("failed to write Feather record batch")?;
328 writer
329 .finish()
330 .context("failed to finalize Feather writer")?;
331
332 Ok(())
333}