1use arrow_array::{Array, ArrayRef, Float64Array, Int64Array, RecordBatch, RecordBatchReader, StringArray};
15use arrow_csv::ReaderBuilder;
16use arrow_schema::{DataType, Field, Schema};
17use lex_bytecode::Value;
18use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
19use parquet::arrow::arrow_writer::ArrowWriter;
20use parquet::arrow::ProjectionMask;
21use parquet::file::properties::WriterProperties;
22use std::collections::VecDeque;
23use std::fs::File;
24use std::path::Path;
25use std::sync::Arc;
26
27fn err<T>(s: impl Into<String>) -> Result<T, String> {
30 Err(s.into())
31}
32
33fn expect_table(v: Option<&Value>) -> Result<&Arc<RecordBatch>, String> {
34 match v {
35 Some(Value::ArrowTable(t)) => Ok(t),
36 Some(other) => err(format!("expected arrow.Table, got {other:?}")),
37 None => err("expected arrow.Table, got nothing"),
38 }
39}
40
41fn expect_str(v: Option<&Value>) -> Result<&str, String> {
42 match v {
43 Some(Value::Str(s)) => Ok(s.as_str()),
44 Some(other) => err(format!("expected Str, got {other:?}")),
45 None => err("expected Str, got nothing"),
46 }
47}
48
49fn expect_int(v: Option<&Value>) -> Result<i64, String> {
50 match v {
51 Some(Value::Int(n)) => Ok(*n),
52 Some(other) => err(format!("expected Int, got {other:?}")),
53 None => err("expected Int, got nothing"),
54 }
55}
56
57fn expect_list(v: Option<&Value>) -> Result<&VecDeque<Value>, String> {
58 match v {
59 Some(Value::List(items)) => Ok(items),
60 Some(other) => err(format!("expected List, got {other:?}")),
61 None => err("expected List, got nothing"),
62 }
63}
64
65fn decode_columns_list(list: &VecDeque<Value>) -> Result<Vec<(&str, &VecDeque<Value>)>, String> {
68 let mut out = Vec::with_capacity(list.len());
69 for (i, item) in list.iter().enumerate() {
70 let pair = match item {
71 Value::Tuple(t) if t.len() == 2 => t,
72 other => {
73 return err(format!(
74 "from_*_columns: column #{i} must be a (Str, List) tuple, got {other:?}"
75 ))
76 }
77 };
78 let name = match &pair[0] {
79 Value::Str(s) => s.as_str(),
80 other => {
81 return err(format!(
82 "from_*_columns: column #{i} name must be Str, got {other:?}"
83 ))
84 }
85 };
86 let values = match &pair[1] {
87 Value::List(items) => items,
88 other => {
89 return err(format!(
90 "from_*_columns: column #{i} (`{name}`) values must be List, got {other:?}"
91 ))
92 }
93 };
94 out.push((name, values));
95 }
96 Ok(out)
97}
98
99fn build_schema_and_check_lengths(cols: &[(&str, ArrayRef)]) -> Result<Schema, String> {
100 if cols.is_empty() {
101 return Ok(Schema::empty());
102 }
103 let nrows = cols[0].1.len();
104 let mut fields = Vec::with_capacity(cols.len());
105 for (name, arr) in cols {
106 if arr.len() != nrows {
107 return err(format!(
108 "from_*_columns: column `{name}` has {} rows, expected {nrows}",
109 arr.len()
110 ));
111 }
112 fields.push(Field::new(*name, arr.data_type().clone(), false));
113 }
114 Ok(Schema::new(fields))
115}
116
117fn pack_table(cols: Vec<(&str, ArrayRef)>) -> Result<Value, String> {
118 let schema = build_schema_and_check_lengths(&cols)?;
119 let arrays: Vec<ArrayRef> = cols.into_iter().map(|(_, a)| a).collect();
120 let batch = RecordBatch::try_new(Arc::new(schema), arrays)
121 .map_err(|e| format!("arrow: failed to build RecordBatch: {e}"))?;
122 Ok(Value::ArrowTable(Arc::new(batch)))
123}
124
125fn from_int_columns(args: &[Value]) -> Result<Value, String> {
129 let list = expect_list(args.first())?;
130 let pairs = decode_columns_list(list)?;
131 let mut owned_names: Vec<String> = Vec::with_capacity(pairs.len());
132 let mut arrays: Vec<ArrayRef> = Vec::with_capacity(pairs.len());
133 for (name, values) in &pairs {
134 owned_names.push((*name).to_string());
135 let mut buf: Vec<i64> = Vec::with_capacity(values.len());
136 for v in values.iter() {
137 match v {
138 Value::Int(n) => buf.push(*n),
139 other => {
140 return err(format!(
141 "from_int_columns: column `{name}` non-Int element: {other:?}"
142 ))
143 }
144 }
145 }
146 arrays.push(Arc::new(Int64Array::from(buf)) as ArrayRef);
147 }
148 let cols: Vec<(&str, ArrayRef)> = owned_names.iter().map(|n| n.as_str()).zip(arrays).collect();
149 pack_table(cols)
150}
151
152fn from_float_columns(args: &[Value]) -> Result<Value, String> {
154 let list = expect_list(args.first())?;
155 let pairs = decode_columns_list(list)?;
156 let mut owned_names: Vec<String> = Vec::with_capacity(pairs.len());
157 let mut arrays: Vec<ArrayRef> = Vec::with_capacity(pairs.len());
158 for (name, values) in &pairs {
159 owned_names.push((*name).to_string());
160 let mut buf: Vec<f64> = Vec::with_capacity(values.len());
161 for v in values.iter() {
162 match v {
163 Value::Float(f) => buf.push(*f),
164 Value::Int(n) => buf.push(*n as f64),
165 other => {
166 return err(format!(
167 "from_float_columns: column `{name}` non-Float element: {other:?}"
168 ))
169 }
170 }
171 }
172 arrays.push(Arc::new(Float64Array::from(buf)) as ArrayRef);
173 }
174 let cols: Vec<(&str, ArrayRef)> = owned_names.iter().map(|n| n.as_str()).zip(arrays).collect();
175 pack_table(cols)
176}
177
178fn from_str_columns(args: &[Value]) -> Result<Value, String> {
180 let list = expect_list(args.first())?;
181 let pairs = decode_columns_list(list)?;
182 let mut owned_names: Vec<String> = Vec::with_capacity(pairs.len());
183 let mut arrays: Vec<ArrayRef> = Vec::with_capacity(pairs.len());
184 for (name, values) in &pairs {
185 owned_names.push((*name).to_string());
186 let mut buf: Vec<String> = Vec::with_capacity(values.len());
187 for v in values.iter() {
188 match v {
189 Value::Str(s) => buf.push(s.to_string()),
190 other => {
191 return err(format!(
192 "from_str_columns: column `{name}` non-Str element: {other:?}"
193 ))
194 }
195 }
196 }
197 arrays.push(Arc::new(StringArray::from(buf)) as ArrayRef);
198 }
199 let cols: Vec<(&str, ArrayRef)> = owned_names.iter().map(|n| n.as_str()).zip(arrays).collect();
200 pack_table(cols)
201}
202
203fn nrows(args: &[Value]) -> Result<Value, String> {
206 Ok(Value::Int(expect_table(args.first())?.num_rows() as i64))
207}
208
209fn ncols(args: &[Value]) -> Result<Value, String> {
210 Ok(Value::Int(expect_table(args.first())?.num_columns() as i64))
211}
212
213fn col_names(args: &[Value]) -> Result<Value, String> {
214 let t = expect_table(args.first())?;
215 let names: VecDeque<Value> = t
216 .schema()
217 .fields()
218 .iter()
219 .map(|f| Value::Str(f.name().as_str().into()))
220 .collect();
221 Ok(Value::List(names))
222}
223
224fn col_type(args: &[Value]) -> Result<Value, String> {
225 let t = expect_table(args.first())?;
226 let name = expect_str(args.get(1))?;
227 match t.schema().column_with_name(name) {
228 None => Ok(none()),
229 Some((_, field)) => Ok(some(Value::Str(format!("{}", field.data_type()).into()))),
230 }
231}
232
233fn lookup_array<'a>(t: &'a RecordBatch, name: &str) -> Result<&'a ArrayRef, String> {
236 let (idx, _) = t
237 .schema()
238 .column_with_name(name)
239 .ok_or_else(|| format!("arrow: column `{name}` not found"))?;
240 Ok(t.column(idx))
241}
242
243fn as_int64<'a>(arr: &'a ArrayRef, name: &str) -> Result<&'a Int64Array, String> {
244 arr.as_any()
245 .downcast_ref::<Int64Array>()
246 .ok_or_else(|| format!("arrow: column `{name}` is {}, not Int64", arr.data_type()))
247}
248
249fn as_float64<'a>(arr: &'a ArrayRef, name: &str) -> Result<&'a Float64Array, String> {
250 arr.as_any()
251 .downcast_ref::<Float64Array>()
252 .ok_or_else(|| format!("arrow: column `{name}` is {}, not Float64", arr.data_type()))
253}
254
255fn col_sum_int(args: &[Value]) -> Result<Value, String> {
256 let t = expect_table(args.first())?;
257 let name = expect_str(args.get(1))?;
258 let arr = as_int64(lookup_array(t, name)?, name)?;
259 let s: i64 = arrow_arith::aggregate::sum(arr).unwrap_or(0);
260 Ok(Value::Int(s))
261}
262
263fn col_sum_float(args: &[Value]) -> Result<Value, String> {
264 let t = expect_table(args.first())?;
265 let name = expect_str(args.get(1))?;
266 let arr = lookup_array(t, name)?;
267 let s = match arr.data_type() {
268 DataType::Float64 => arrow_arith::aggregate::sum(as_float64(arr, name)?).unwrap_or(0.0),
269 DataType::Int64 => arrow_arith::aggregate::sum(as_int64(arr, name)?).unwrap_or(0) as f64,
270 other => {
271 return err(format!(
272 "col_sum_float: column `{name}` is {other:?}, expected Int64 or Float64"
273 ))
274 }
275 };
276 Ok(Value::Float(s))
277}
278
279fn col_mean(args: &[Value]) -> Result<Value, String> {
280 let t = expect_table(args.first())?;
281 let name = expect_str(args.get(1))?;
282 let arr = lookup_array(t, name)?;
283 let n = arr.len() as f64 - arr.null_count() as f64;
284 if n == 0.0 {
285 return Ok(none());
286 }
287 let total: f64 = match arr.data_type() {
288 DataType::Float64 => arrow_arith::aggregate::sum(as_float64(arr, name)?).unwrap_or(0.0),
289 DataType::Int64 => arrow_arith::aggregate::sum(as_int64(arr, name)?).unwrap_or(0) as f64,
290 other => {
291 return err(format!(
292 "col_mean: column `{name}` is {other:?}, expected Int64 or Float64"
293 ))
294 }
295 };
296 Ok(some(Value::Float(total / n)))
297}
298
299fn col_min_int(args: &[Value]) -> Result<Value, String> {
300 let t = expect_table(args.first())?;
301 let name = expect_str(args.get(1))?;
302 let arr = as_int64(lookup_array(t, name)?, name)?;
303 match arrow_arith::aggregate::min(arr) {
304 Some(v) => Ok(some(Value::Int(v))),
305 None => Ok(none()),
306 }
307}
308
309fn col_max_int(args: &[Value]) -> Result<Value, String> {
310 let t = expect_table(args.first())?;
311 let name = expect_str(args.get(1))?;
312 let arr = as_int64(lookup_array(t, name)?, name)?;
313 match arrow_arith::aggregate::max(arr) {
314 Some(v) => Ok(some(Value::Int(v))),
315 None => Ok(none()),
316 }
317}
318
319fn col_count(args: &[Value]) -> Result<Value, String> {
320 let t = expect_table(args.first())?;
321 let name = expect_str(args.get(1))?;
322 let arr = lookup_array(t, name)?;
323 Ok(Value::Int((arr.len() - arr.null_count()) as i64))
324}
325
326fn head(args: &[Value]) -> Result<Value, String> {
329 let t = expect_table(args.first())?;
330 let n = expect_int(args.get(1))?.max(0) as usize;
331 let take = n.min(t.num_rows());
332 Ok(Value::ArrowTable(Arc::new(t.slice(0, take))))
333}
334
335fn tail(args: &[Value]) -> Result<Value, String> {
336 let t = expect_table(args.first())?;
337 let n = expect_int(args.get(1))?.max(0) as usize;
338 let total = t.num_rows();
339 let take = n.min(total);
340 Ok(Value::ArrowTable(Arc::new(t.slice(total - take, take))))
341}
342
343fn slice(args: &[Value]) -> Result<Value, String> {
344 let t = expect_table(args.first())?;
345 let start = expect_int(args.get(1))?.max(0) as usize;
346 let stop = expect_int(args.get(2))?.max(0) as usize;
347 let total = t.num_rows();
348 let s = start.min(total);
349 let e = stop.min(total).max(s);
350 Ok(Value::ArrowTable(Arc::new(t.slice(s, e - s))))
351}
352
353fn select_cols(args: &[Value]) -> Result<Value, String> {
354 let t = expect_table(args.first())?;
355 let names_list = expect_list(args.get(1))?;
356 let mut indices = Vec::with_capacity(names_list.len());
357 for v in names_list.iter() {
358 let n = match v {
359 Value::Str(s) => s.as_str(),
360 other => {
361 return err(format!(
362 "select_cols: name list contained non-Str: {other:?}"
363 ))
364 }
365 };
366 let (i, _) = t
367 .schema()
368 .column_with_name(n)
369 .ok_or_else(|| format!("select_cols: column `{n}` not found"))?;
370 indices.push(i);
371 }
372 let projected = t
373 .project(&indices)
374 .map_err(|e| format!("select_cols: {e}"))?;
375 Ok(Value::ArrowTable(Arc::new(projected)))
376}
377
378fn rename_col(args: &[Value]) -> Result<Value, String> {
379 let t = expect_table(args.first())?;
380 let old_name = expect_str(args.get(1))?;
381 let new_name = expect_str(args.get(2))?;
382 let (idx, _) = t
383 .schema()
384 .column_with_name(old_name)
385 .ok_or_else(|| format!("rename_col: column `{old_name}` not found"))?;
386 let fields: Vec<Field> = t
387 .schema()
388 .fields()
389 .iter()
390 .enumerate()
391 .map(|(i, f)| {
392 if i == idx {
393 Field::new(new_name, f.data_type().clone(), f.is_nullable())
394 } else {
395 f.as_ref().clone()
396 }
397 })
398 .collect();
399 let renamed = RecordBatch::try_new(Arc::new(Schema::new(fields)), t.columns().to_vec())
400 .map_err(|e| format!("rename_col: {e}"))?;
401 Ok(Value::ArrowTable(Arc::new(renamed)))
402}
403
404fn drop_col(args: &[Value]) -> Result<Value, String> {
405 let t = expect_table(args.first())?;
406 let drop_name = expect_str(args.get(1))?;
407 let mut keep = Vec::with_capacity(t.num_columns());
408 for (i, f) in t.schema().fields().iter().enumerate() {
409 if f.name() != drop_name {
410 keep.push(i);
411 }
412 }
413 if keep.len() == t.num_columns() {
414 return err(format!("drop_col: column `{drop_name}` not found"));
415 }
416 let projected = t.project(&keep).map_err(|e| format!("drop_col: {e}"))?;
417 Ok(Value::ArrowTable(Arc::new(projected)))
418}
419
420pub fn read_csv_at(path: &Path) -> Result<Value, String> {
440 #[cfg(feature = "df")]
441 {
442 crate::df::read_csv_at_polars(path)
443 }
444 #[cfg(not(feature = "df"))]
445 {
446 read_csv_at_arrow_rs(path)
447 }
448}
449
450#[cfg_attr(feature = "df", allow(dead_code))]
451fn read_csv_at_arrow_rs(path: &Path) -> Result<Value, String> {
452 let file =
453 File::open(path).map_err(|e| format!("arrow.read_csv: open `{}`: {e}", path.display()))?;
454 let (schema, _) = arrow_csv::reader::Format::default()
455 .with_header(true)
456 .infer_schema(&file, Some(100))
457 .map_err(|e| format!("arrow.read_csv: schema inference: {e}"))?;
458 let file = File::open(path)
460 .map_err(|e| format!("arrow.read_csv: reopen `{}`: {e}", path.display()))?;
461 let schema = Arc::new(schema);
462 let reader = ReaderBuilder::new(Arc::clone(&schema))
463 .with_header(true)
464 .build(file)
465 .map_err(|e| format!("arrow.read_csv: reader build: {e}"))?;
466 let mut batches: Vec<RecordBatch> = Vec::new();
467 for batch in reader {
468 batches.push(batch.map_err(|e| format!("arrow.read_csv: row decode: {e}"))?);
469 }
470 let combined = if batches.is_empty() {
471 RecordBatch::new_empty(schema)
472 } else {
473 arrow_select::concat::concat_batches(&schema, &batches)
474 .map_err(|e| format!("arrow.read_csv: concat: {e}"))?
475 };
476 Ok(Value::ArrowTable(Arc::new(combined)))
477}
478
479pub fn read_parquet_at(path: &Path) -> Result<Value, String> {
488 let file = File::open(path)
489 .map_err(|e| format!("arrow.read_parquet: open `{}`: {e}", path.display()))?;
490 let builder = ParquetRecordBatchReaderBuilder::try_new(file)
491 .map_err(|e| format!("arrow.read_parquet: open `{}`: {e}", path.display()))?;
492 let schema = builder.schema().clone();
493 let reader = builder
494 .build()
495 .map_err(|e| format!("arrow.read_parquet: reader build: {e}"))?;
496 let mut batches: Vec<RecordBatch> = Vec::new();
497 for batch in reader {
498 batches.push(batch.map_err(|e| format!("arrow.read_parquet: row-group decode: {e}"))?);
499 }
500 let combined = if batches.is_empty() {
501 RecordBatch::new_empty(schema)
502 } else {
503 arrow_select::concat::concat_batches(&schema, &batches)
504 .map_err(|e| format!("arrow.read_parquet: concat: {e}"))?
505 };
506 Ok(Value::ArrowTable(Arc::new(combined)))
507}
508
509pub fn read_parquet_cols_at(path: &Path, cols: &[String]) -> Result<Value, String> {
513 let file = File::open(path)
514 .map_err(|e| format!("arrow.read_parquet_cols: open `{}`: {e}", path.display()))?;
515 let builder = ParquetRecordBatchReaderBuilder::try_new(file)
516 .map_err(|e| format!("arrow.read_parquet_cols: open `{}`: {e}", path.display()))?;
517 let mask = {
522 let parquet_schema = builder.parquet_schema();
523 let root_fields = parquet_schema.root_schema().get_fields();
524 let mut indices = Vec::with_capacity(cols.len());
525 for name in cols {
526 let idx = root_fields
527 .iter()
528 .position(|f| f.name() == name.as_str())
529 .ok_or_else(|| format!("arrow.read_parquet_cols: column `{name}` not in file"))?;
530 indices.push(idx);
531 }
532 ProjectionMask::roots(parquet_schema, indices)
533 };
534 let reader = builder
535 .with_projection(mask)
536 .build()
537 .map_err(|e| format!("arrow.read_parquet_cols: reader build: {e}"))?;
538 let projected_schema = reader.schema();
539 let mut batches: Vec<RecordBatch> = Vec::new();
540 for batch in reader {
541 batches
542 .push(batch.map_err(|e| format!("arrow.read_parquet_cols: row-group decode: {e}"))?);
543 }
544 let combined = if batches.is_empty() {
545 RecordBatch::new_empty(projected_schema)
546 } else {
547 arrow_select::concat::concat_batches(&projected_schema, &batches)
548 .map_err(|e| format!("arrow.read_parquet_cols: concat: {e}"))?
549 };
550 Ok(Value::ArrowTable(Arc::new(combined)))
551}
552
553pub fn write_parquet_at(rb: &RecordBatch, path: &Path) -> Result<Value, String> {
556 let file = File::create(path)
557 .map_err(|e| format!("arrow.write_parquet: create `{}`: {e}", path.display()))?;
558 let props = WriterProperties::builder().build();
559 let mut writer = ArrowWriter::try_new(file, rb.schema(), Some(props))
560 .map_err(|e| format!("arrow.write_parquet: writer init: {e}"))?;
561 writer
562 .write(rb)
563 .map_err(|e| format!("arrow.write_parquet: write: {e}"))?;
564 writer
565 .close()
566 .map_err(|e| format!("arrow.write_parquet: close: {e}"))?;
567 Ok(Value::Unit)
568}
569
570pub fn write_csv_at(rb: &RecordBatch, path: &Path) -> Result<Value, String> {
573 let file = File::create(path)
574 .map_err(|e| format!("arrow.write_csv: create `{}`: {e}", path.display()))?;
575 let mut writer = arrow_csv::WriterBuilder::new()
576 .with_header(true)
577 .build(file);
578 writer
579 .write(rb)
580 .map_err(|e| format!("arrow.write_csv: write: {e}"))?;
581 Ok(Value::Unit)
582}
583
584fn some(v: Value) -> Value {
587 Value::Variant {
588 name: "Some".into(),
589 args: vec![v],
590 }
591}
592
593fn none() -> Value {
594 Value::Variant {
595 name: "None".into(),
596 args: vec![],
597 }
598}
599
600fn ok(v: Value) -> Value {
601 Value::Variant {
602 name: "Ok".into(),
603 args: vec![v],
604 }
605}
606
607fn err_variant(s: String) -> Value {
608 Value::Variant {
609 name: "Err".into(),
610 args: vec![Value::Str(s.into())],
611 }
612}
613
614fn lift_result(r: Result<Value, String>) -> Result<Value, String> {
620 match r {
621 Ok(v) => Ok(ok(v)),
622 Err(s) => Ok(err_variant(s)),
623 }
624}
625
626pub fn dispatch(op: &str, args: &[Value]) -> Option<Result<Value, String>> {
638 Some(match op {
639 "from_int_columns" => lift_result(from_int_columns(args)),
641 "from_float_columns" => lift_result(from_float_columns(args)),
642 "from_str_columns" => lift_result(from_str_columns(args)),
643 "col_sum_int" => lift_result(col_sum_int(args)),
644 "col_sum_float" => lift_result(col_sum_float(args)),
645 "col_mean" => lift_result(col_mean(args)),
646 "col_min_int" => lift_result(col_min_int(args)),
647 "col_max_int" => lift_result(col_max_int(args)),
648 "col_count" => lift_result(col_count(args)),
649 "select_cols" => lift_result(select_cols(args)),
650 "drop_col" => lift_result(drop_col(args)),
651 "rename_col" => lift_result(rename_col(args)),
652 "nrows" => nrows(args),
654 "ncols" => ncols(args),
655 "col_names" => col_names(args),
656 "col_type" => col_type(args),
657 "head" => head(args),
658 "tail" => tail(args),
659 "slice" => slice(args),
660 _ => return None,
661 })
662}