Skip to content

Commit db0ab74

Browse files
cj-zhukovSergey Zhukov
cj-zhukov
and
Sergey Zhukov
authored
Add macro for creating DataFrame (#16090) (#16104)
* Add macro for creating DataFrame (#16090) --------- Co-authored-by: Sergey Zhukov <[email protected]>
1 parent 260a28a commit db0ab74

File tree

6 files changed

+546
-2
lines changed

6 files changed

+546
-2
lines changed

‎datafusion-examples/examples/dataframe.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ async fn main() -> Result<()> {
6363
read_parquet(&ctx).await?;
6464
read_csv(&ctx).await?;
6565
read_memory(&ctx).await?;
66+
read_memory_macro().await?;
6667
write_out(&ctx).await?;
6768
register_aggregate_test_data("t1", &ctx).await?;
6869
register_aggregate_test_data("t2", &ctx).await?;
@@ -173,6 +174,24 @@ async fn read_memory(ctx: &SessionContext) -> Result<()> {
173174
Ok(())
174175
}
175176

177+
/// Use the DataFrame API to:
178+
/// 1. Read in-memory data.
179+
async fn read_memory_macro() -> Result<()> {
180+
// create a DataFrame using macro
181+
let df = dataframe!(
182+
"a" => ["a", "b", "c", "d"],
183+
"b" => [1, 10, 10, 100]
184+
)?;
185+
// print the results
186+
df.show().await?;
187+
188+
// create empty DataFrame using macro
189+
let df_empty = dataframe!()?;
190+
df_empty.show().await?;
191+
192+
Ok(())
193+
}
194+
176195
/// Use the DataFrame API to:
177196
/// 1. Write out a DataFrame to a table
178197
/// 2. Write out a DataFrame to a parquet file

‎datafusion/common/src/test_util.rs‎

Lines changed: 326 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,16 @@
1818
//! Utility functions to make testing DataFusion based crates easier
1919
2020
use crate::arrow::util::pretty::pretty_format_batches_with_options;
21-
use arrow::array::RecordBatch;
21+
use arrow::array::{ArrayRef, RecordBatch};
2222
use arrow::error::ArrowError;
2323
use std::fmt::Display;
2424
use std::{error::Error, path::PathBuf};
2525

26+
/// Converts a vector or array into an ArrayRef.
27+
pub trait IntoArrayRef {
28+
fn into_array_ref(self) -> ArrayRef;
29+
}
30+
2631
pub fn format_batches(results: &[RecordBatch]) -> Result<impl Display, ArrowError> {
2732
let datafusion_format_options = crate::config::FormatOptions::default();
2833

@@ -383,6 +388,326 @@ macro_rules! record_batch {
383388
}
384389
}
385390

391+
pub mod array_conversion {
392+
use arrow::array::ArrayRef;
393+
394+
use super::IntoArrayRef;
395+
396+
impl IntoArrayRef for Vec<bool> {
397+
fn into_array_ref(self) -> ArrayRef {
398+
create_array!(Boolean, self)
399+
}
400+
}
401+
402+
impl IntoArrayRef for Vec<Option<bool>> {
403+
fn into_array_ref(self) -> ArrayRef {
404+
create_array!(Boolean, self)
405+
}
406+
}
407+
408+
impl IntoArrayRef for &[bool] {
409+
fn into_array_ref(self) -> ArrayRef {
410+
create_array!(Boolean, self.to_vec())
411+
}
412+
}
413+
414+
impl IntoArrayRef for &[Option<bool>] {
415+
fn into_array_ref(self) -> ArrayRef {
416+
create_array!(Boolean, self.to_vec())
417+
}
418+
}
419+
420+
impl IntoArrayRef for Vec<i8> {
421+
fn into_array_ref(self) -> ArrayRef {
422+
create_array!(Int8, self)
423+
}
424+
}
425+
426+
impl IntoArrayRef for Vec<Option<i8>> {
427+
fn into_array_ref(self) -> ArrayRef {
428+
create_array!(Int8, self)
429+
}
430+
}
431+
432+
impl IntoArrayRef for &[i8] {
433+
fn into_array_ref(self) -> ArrayRef {
434+
create_array!(Int8, self.to_vec())
435+
}
436+
}
437+
438+
impl IntoArrayRef for &[Option<i8>] {
439+
fn into_array_ref(self) -> ArrayRef {
440+
create_array!(Int8, self.to_vec())
441+
}
442+
}
443+
444+
impl IntoArrayRef for Vec<i16> {
445+
fn into_array_ref(self) -> ArrayRef {
446+
create_array!(Int16, self)
447+
}
448+
}
449+
450+
impl IntoArrayRef for Vec<Option<i16>> {
451+
fn into_array_ref(self) -> ArrayRef {
452+
create_array!(Int16, self)
453+
}
454+
}
455+
456+
impl IntoArrayRef for &[i16] {
457+
fn into_array_ref(self) -> ArrayRef {
458+
create_array!(Int16, self.to_vec())
459+
}
460+
}
461+
462+
impl IntoArrayRef for &[Option<i16>] {
463+
fn into_array_ref(self) -> ArrayRef {
464+
create_array!(Int16, self.to_vec())
465+
}
466+
}
467+
468+
impl IntoArrayRef for Vec<i32> {
469+
fn into_array_ref(self) -> ArrayRef {
470+
create_array!(Int32, self)
471+
}
472+
}
473+
474+
impl IntoArrayRef for Vec<Option<i32>> {
475+
fn into_array_ref(self) -> ArrayRef {
476+
create_array!(Int32, self)
477+
}
478+
}
479+
480+
impl IntoArrayRef for &[i32] {
481+
fn into_array_ref(self) -> ArrayRef {
482+
create_array!(Int32, self.to_vec())
483+
}
484+
}
485+
486+
impl IntoArrayRef for &[Option<i32>] {
487+
fn into_array_ref(self) -> ArrayRef {
488+
create_array!(Int32, self.to_vec())
489+
}
490+
}
491+
492+
impl IntoArrayRef for Vec<i64> {
493+
fn into_array_ref(self) -> ArrayRef {
494+
create_array!(Int64, self)
495+
}
496+
}
497+
498+
impl IntoArrayRef for Vec<Option<i64>> {
499+
fn into_array_ref(self) -> ArrayRef {
500+
create_array!(Int64, self)
501+
}
502+
}
503+
504+
impl IntoArrayRef for &[i64] {
505+
fn into_array_ref(self) -> ArrayRef {
506+
create_array!(Int64, self.to_vec())
507+
}
508+
}
509+
510+
impl IntoArrayRef for &[Option<i64>] {
511+
fn into_array_ref(self) -> ArrayRef {
512+
create_array!(Int64, self.to_vec())
513+
}
514+
}
515+
516+
impl IntoArrayRef for Vec<u8> {
517+
fn into_array_ref(self) -> ArrayRef {
518+
create_array!(UInt8, self)
519+
}
520+
}
521+
522+
impl IntoArrayRef for Vec<Option<u8>> {
523+
fn into_array_ref(self) -> ArrayRef {
524+
create_array!(UInt8, self)
525+
}
526+
}
527+
528+
impl IntoArrayRef for &[u8] {
529+
fn into_array_ref(self) -> ArrayRef {
530+
create_array!(UInt8, self.to_vec())
531+
}
532+
}
533+
534+
impl IntoArrayRef for &[Option<u8>] {
535+
fn into_array_ref(self) -> ArrayRef {
536+
create_array!(UInt8, self.to_vec())
537+
}
538+
}
539+
540+
impl IntoArrayRef for Vec<u16> {
541+
fn into_array_ref(self) -> ArrayRef {
542+
create_array!(UInt16, self)
543+
}
544+
}
545+
546+
impl IntoArrayRef for Vec<Option<u16>> {
547+
fn into_array_ref(self) -> ArrayRef {
548+
create_array!(UInt16, self)
549+
}
550+
}
551+
552+
impl IntoArrayRef for &[u16] {
553+
fn into_array_ref(self) -> ArrayRef {
554+
create_array!(UInt16, self.to_vec())
555+
}
556+
}
557+
558+
impl IntoArrayRef for &[Option<u16>] {
559+
fn into_array_ref(self) -> ArrayRef {
560+
create_array!(UInt16, self.to_vec())
561+
}
562+
}
563+
564+
impl IntoArrayRef for Vec<u32> {
565+
fn into_array_ref(self) -> ArrayRef {
566+
create_array!(UInt32, self)
567+
}
568+
}
569+
570+
impl IntoArrayRef for Vec<Option<u32>> {
571+
fn into_array_ref(self) -> ArrayRef {
572+
create_array!(UInt32, self)
573+
}
574+
}
575+
576+
impl IntoArrayRef for &[u32] {
577+
fn into_array_ref(self) -> ArrayRef {
578+
create_array!(UInt32, self.to_vec())
579+
}
580+
}
581+
582+
impl IntoArrayRef for &[Option<u32>] {
583+
fn into_array_ref(self) -> ArrayRef {
584+
create_array!(UInt32, self.to_vec())
585+
}
586+
}
587+
588+
impl IntoArrayRef for Vec<u64> {
589+
fn into_array_ref(self) -> ArrayRef {
590+
create_array!(UInt64, self)
591+
}
592+
}
593+
594+
impl IntoArrayRef for Vec<Option<u64>> {
595+
fn into_array_ref(self) -> ArrayRef {
596+
create_array!(UInt64, self)
597+
}
598+
}
599+
600+
impl IntoArrayRef for &[u64] {
601+
fn into_array_ref(self) -> ArrayRef {
602+
create_array!(UInt64, self.to_vec())
603+
}
604+
}
605+
606+
impl IntoArrayRef for &[Option<u64>] {
607+
fn into_array_ref(self) -> ArrayRef {
608+
create_array!(UInt64, self.to_vec())
609+
}
610+
}
611+
612+
//#TODO add impl for f16
613+
614+
impl IntoArrayRef for Vec<f32> {
615+
fn into_array_ref(self) -> ArrayRef {
616+
create_array!(Float32, self)
617+
}
618+
}
619+
620+
impl IntoArrayRef for Vec<Option<f32>> {
621+
fn into_array_ref(self) -> ArrayRef {
622+
create_array!(Float32, self)
623+
}
624+
}
625+
626+
impl IntoArrayRef for &[f32] {
627+
fn into_array_ref(self) -> ArrayRef {
628+
create_array!(Float32, self.to_vec())
629+
}
630+
}
631+
632+
impl IntoArrayRef for &[Option<f32>] {
633+
fn into_array_ref(self) -> ArrayRef {
634+
create_array!(Float32, self.to_vec())
635+
}
636+
}
637+
638+
impl IntoArrayRef for Vec<f64> {
639+
fn into_array_ref(self) -> ArrayRef {
640+
create_array!(Float64, self)
641+
}
642+
}
643+
644+
impl IntoArrayRef for Vec<Option<f64>> {
645+
fn into_array_ref(self) -> ArrayRef {
646+
create_array!(Float64, self)
647+
}
648+
}
649+
650+
impl IntoArrayRef for &[f64] {
651+
fn into_array_ref(self) -> ArrayRef {
652+
create_array!(Float64, self.to_vec())
653+
}
654+
}
655+
656+
impl IntoArrayRef for &[Option<f64>] {
657+
fn into_array_ref(self) -> ArrayRef {
658+
create_array!(Float64, self.to_vec())
659+
}
660+
}
661+
662+
impl IntoArrayRef for Vec<&str> {
663+
fn into_array_ref(self) -> ArrayRef {
664+
create_array!(Utf8, self)
665+
}
666+
}
667+
668+
impl IntoArrayRef for Vec<Option<&str>> {
669+
fn into_array_ref(self) -> ArrayRef {
670+
create_array!(Utf8, self)
671+
}
672+
}
673+
674+
impl IntoArrayRef for &[&str] {
675+
fn into_array_ref(self) -> ArrayRef {
676+
create_array!(Utf8, self.to_vec())
677+
}
678+
}
679+
680+
impl IntoArrayRef for &[Option<&str>] {
681+
fn into_array_ref(self) -> ArrayRef {
682+
create_array!(Utf8, self.to_vec())
683+
}
684+
}
685+
686+
impl IntoArrayRef for Vec<String> {
687+
fn into_array_ref(self) -> ArrayRef {
688+
create_array!(Utf8, self)
689+
}
690+
}
691+
692+
impl IntoArrayRef for Vec<Option<String>> {
693+
fn into_array_ref(self) -> ArrayRef {
694+
create_array!(Utf8, self)
695+
}
696+
}
697+
698+
impl IntoArrayRef for &[String] {
699+
fn into_array_ref(self) -> ArrayRef {
700+
create_array!(Utf8, self.to_vec())
701+
}
702+
}
703+
704+
impl IntoArrayRef for &[Option<String>] {
705+
fn into_array_ref(self) -> ArrayRef {
706+
create_array!(Utf8, self.to_vec())
707+
}
708+
}
709+
}
710+
386711
#[cfg(test)]
387712
mod tests {
388713
use crate::cast::{as_float64_array, as_int32_array, as_string_array};

0 commit comments

Comments
 (0)