1use arrow::array::{Array, as_boolean_array, as_primitive_array, as_string_array};
2use arrow::datatypes::*;
3use arrow::record_batch::RecordBatch;
4use chrono::{DateTime, NaiveDate, Utc};
5use rust_xlsxwriter::{Workbook, XlsxError};
6
7#[derive(Debug)]
8pub enum Error {
9 Xlsx(XlsxError),
10 Arrow(arrow::error::ArrowError),
11}
12
13impl From<XlsxError> for Error {
14 fn from(e: XlsxError) -> Self {
15 Error::Xlsx(e)
16 }
17}
18
19impl From<arrow::error::ArrowError> for Error {
20 fn from(e: arrow::error::ArrowError) -> Self {
21 Error::Arrow(e)
22 }
23}
24
25impl std::fmt::Display for Error {
26 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27 match self {
28 Error::Xlsx(e) => write!(f, "Xlsx error: {}", e),
29 Error::Arrow(e) => write!(f, "Arrow error: {}", e),
30 }
31 }
32}
33
34impl std::error::Error for Error {}
35
36pub fn batch_iter2x<I>(mut bi: I, book: &mut Workbook, sheet_name: &str) -> Result<(), Error>
37where
38 I: Iterator<Item = Result<RecordBatch, arrow::error::ArrowError>>,
39{
40 let worksheet = book.add_worksheet().set_name(sheet_name)?;
41
42 let mut row_offset = 0;
43
44 if let Some(batch_result) = bi.next() {
45 let batch = batch_result?;
46 let schema = batch.schema();
47 for (col, field) in schema.fields().iter().enumerate() {
48 worksheet.write_string(row_offset, col as u16, field.name())?;
49 }
50 row_offset += 1;
51
52 write_batch(worksheet, &batch, &mut row_offset)?;
53
54 for batch_result in bi {
55 let batch = batch_result?;
56 write_batch(worksheet, &batch, &mut row_offset)?;
57 }
58 }
59
60 Ok(())
61}
62
63macro_rules! write_primitive_number {
64 ($worksheet:expr, $column:expr, $row:expr, $row_offset:expr, $col:expr, $type:ty) => {{
65 let array = as_primitive_array::<$type>($column);
66 if !array.is_null($row) {
67 let value = array.value($row);
68 $worksheet.write_number($row_offset, $col as u16, value as f64)?;
69 }
70 }};
71}
72
73fn write_batch(
74 worksheet: &mut rust_xlsxwriter::Worksheet,
75 batch: &RecordBatch,
76 row_offset: &mut u32,
77) -> Result<(), XlsxError> {
78 for row in 0..batch.num_rows() {
79 for col in 0..batch.num_columns() {
80 let column = batch.column(col);
81 let data_type = column.data_type();
82
83 match data_type {
84 DataType::Utf8 => {
85 let array = as_string_array(column);
86 if !array.is_null(row) {
87 let value = array.value(row);
88 worksheet.write_string(*row_offset, col as u16, value)?;
89 }
90 }
91 DataType::Int8 => {
92 write_primitive_number!(worksheet, column, row, *row_offset, col, Int8Type)
93 }
94 DataType::Int16 => {
95 write_primitive_number!(worksheet, column, row, *row_offset, col, Int16Type)
96 }
97 DataType::Int32 => {
98 write_primitive_number!(worksheet, column, row, *row_offset, col, Int32Type)
99 }
100 DataType::Int64 => {
101 write_primitive_number!(worksheet, column, row, *row_offset, col, Int64Type)
102 }
103 DataType::UInt8 => {
104 write_primitive_number!(worksheet, column, row, *row_offset, col, UInt8Type)
105 }
106 DataType::UInt16 => {
107 write_primitive_number!(worksheet, column, row, *row_offset, col, UInt16Type)
108 }
109 DataType::UInt32 => {
110 write_primitive_number!(worksheet, column, row, *row_offset, col, UInt32Type)
111 }
112 DataType::UInt64 => {
113 write_primitive_number!(worksheet, column, row, *row_offset, col, UInt64Type)
114 }
115 DataType::Float16 => {
116 let array = as_primitive_array::<Float16Type>(column);
117 if !array.is_null(row) {
118 let value = array.value(row);
119 worksheet.write_number(*row_offset, col as u16, value.to_f64())?;
120 }
121 }
122 DataType::Float32 => {
123 write_primitive_number!(worksheet, column, row, *row_offset, col, Float32Type)
124 }
125 DataType::Float64 => {
126 write_primitive_number!(worksheet, column, row, *row_offset, col, Float64Type)
127 }
128 DataType::Boolean => {
129 let array = as_boolean_array(column);
130 if !array.is_null(row) {
131 let value = array.value(row);
132 worksheet.write_boolean(*row_offset, col as u16, value)?;
133 }
134 }
135 DataType::Date32 => {
136 let array = as_primitive_array::<Date32Type>(column);
137 if !array.is_null(row) {
138 let value = array.value(row);
139 if let Some(date) = NaiveDate::from_epoch_days(value) {
140 worksheet.write_datetime(*row_offset, col as u16, date)?;
141 }
142 }
143 }
144 DataType::Date64 => {
145 let array = as_primitive_array::<Date64Type>(column);
146 if !array.is_null(row) {
147 let value = array.value(row);
148 if let Some(datetime_utc) = DateTime::<Utc>::from_timestamp_millis(value) {
149 worksheet.write_datetime(
150 *row_offset,
151 col as u16,
152 datetime_utc.naive_utc(),
153 )?;
154 }
155 }
156 }
157 DataType::Time32(unit) => match unit {
158 TimeUnit::Second => {
159 let array = as_primitive_array::<Time32SecondType>(column);
160 if !array.is_null(row) {
161 let value = array.value(row);
162 worksheet.write_number(
163 *row_offset,
164 col as u16,
165 value as f64 / 86400.0,
166 )?;
167 }
168 }
169 TimeUnit::Millisecond => {
170 let array = as_primitive_array::<Time32MillisecondType>(column);
171 if !array.is_null(row) {
172 let value = array.value(row);
173 worksheet.write_number(
174 *row_offset,
175 col as u16,
176 value as f64 / 86_400_000.0,
177 )?;
178 }
179 }
180 _ => {}
181 },
182 DataType::Time64(unit) => match unit {
183 TimeUnit::Microsecond => {
184 let array = as_primitive_array::<Time64MicrosecondType>(column);
185 if !array.is_null(row) {
186 let value = array.value(row);
187 worksheet.write_number(
188 *row_offset,
189 col as u16,
190 value as f64 / 86_400_000_000.0,
191 )?;
192 }
193 }
194 TimeUnit::Nanosecond => {
195 let array = as_primitive_array::<Time64NanosecondType>(column);
196 if !array.is_null(row) {
197 let value = array.value(row);
198 worksheet.write_number(
199 *row_offset,
200 col as u16,
201 value as f64 / 86_400_000_000_000.0,
202 )?;
203 }
204 }
205 _ => {}
206 },
207 DataType::Timestamp(unit, _) => {
208 let s = match unit {
209 TimeUnit::Second => {
210 let array =
211 as_primitive_array::<arrow::datatypes::TimestampSecondType>(column);
212 if !array.is_null(row) {
213 format!("{}", array.value(row))
214 } else {
215 String::new()
216 }
217 }
218 TimeUnit::Millisecond => {
219 let array = as_primitive_array::<
220 arrow::datatypes::TimestampMillisecondType,
221 >(column);
222 if !array.is_null(row) {
223 format!("{}", array.value(row))
224 } else {
225 String::new()
226 }
227 }
228 TimeUnit::Microsecond => {
229 let array = as_primitive_array::<
230 arrow::datatypes::TimestampMicrosecondType,
231 >(column);
232 if !array.is_null(row) {
233 format!("{}", array.value(row))
234 } else {
235 String::new()
236 }
237 }
238 TimeUnit::Nanosecond => {
239 let array = as_primitive_array::<
240 arrow::datatypes::TimestampNanosecondType,
241 >(column);
242 if !array.is_null(row) {
243 format!("{}", array.value(row))
244 } else {
245 String::new()
246 }
247 }
248 };
249 worksheet.write_string(*row_offset, col as u16, s)?;
250 }
251 _ => {
252 worksheet.write_string(
253 *row_offset,
254 col as u16,
255 format!("unsupported data type: {:?}", data_type),
256 )?;
257 }
258 }
259 }
260 *row_offset += 1;
261 }
262 Ok(())
263}