Merge branch 'main' into transform-with-payload

.github/workflows/rust.yml

@@ -99,6 +99,14 @@ jobs:
           rust-version: stable
       - name: Run tests (excluding doctests)
         run: cargo test --lib --tests --bins --features avro,json,backtrace
+        env:
+          # do not produce debug symbols to keep memory usage down
+          # hardcoding other profile params to avoid profile override values
+          # More on Cargo profiles https://doc.rust-lang.org/cargo/reference/profiles.html?profile-settings#profile-settings
+          RUSTFLAGS: "-C debuginfo=0 -C opt-level=0 -C incremental=false -C codegen-units=256"
+          RUST_BACKTRACE: "1"
+          # avoid rust stack overflows on tpc-ds tests
+          RUST_MIN_STACK: "3000000"
       - name: Verify Working Directory Clean
         run: git diff --exit-code
 
@@ -290,6 +298,7 @@ jobs:
       # with a OS-dependent path.
       - name: Setup Rust toolchain
         run: |
+          rustup update stable
           rustup toolchain install stable
           rustup default stable
           rustup component add rustfmt
@@ -301,10 +310,13 @@ jobs:
           cd datafusion-cli
           cargo test --lib --tests --bins --all-features
         env:
-          # do not produce debug symbols to keep memory usage down
-          RUSTFLAGS: "-C debuginfo=0"
+          # Minimize producing debug symbols to keep memory usage down
+          # Set debuginfo=line-tables-only as debuginfo=0 causes immensely slow build
+          # See for more details: https://github.com/rust-lang/rust/issues/119560
+          RUSTFLAGS: "-C debuginfo=line-tables-only"
           RUST_BACKTRACE: "1"
-
+          # avoid rust stack overflows on tpc-ds tests
+          RUST_MIN_STACK: "3000000"
   macos:
     name: cargo test (mac)
     runs-on: macos-latest
@@ -327,6 +339,7 @@ jobs:
       # with a OS-dependent path.
       - name: Setup Rust toolchain
         run: |
+          rustup update stable
           rustup toolchain install stable
           rustup default stable
           rustup component add rustfmt
@@ -338,8 +351,12 @@ jobs:
           cargo test --lib --tests --bins --all-features
         env:
           # do not produce debug symbols to keep memory usage down
-          RUSTFLAGS: "-C debuginfo=0"
+          # hardcoding other profile params to avoid profile override values
+          # More on Cargo profiles https://doc.rust-lang.org/cargo/reference/profiles.html?profile-settings#profile-settings
+          RUSTFLAGS: "-C debuginfo=0 -C opt-level=0 -C incremental=false -C codegen-units=256"
           RUST_BACKTRACE: "1"
+          # avoid rust stack overflows on tpc-ds tests
+          RUST_MIN_STACK: "3000000"
 
   test-datafusion-pyarrow:
     name: cargo test pyarrow (amd64)

datafusion-cli/Cargo.toml

@@ -45,7 +45,7 @@ parking_lot = { version = "0.12" }
 parquet = { version = "49.0.0", default-features = false }
 regex = "1.8"
 rustyline = "11.0"
-tokio = { version = "1.24", features = ["macros", "rt", "rt-multi-thread", "sync", "parking_lot"] }
+tokio = { version = "1.24", features = ["macros", "rt", "rt-multi-thread", "sync", "parking_lot", "signal"] }
 url = "2.2"
 
 [dev-dependencies]

datafusion-cli/src/exec.rs

@@ -45,13 +45,14 @@ use datafusion::sql::{parser::DFParser, sqlparser::dialect::dialect_from_str};
 use object_store::ObjectStore;
 use rustyline::error::ReadlineError;
 use rustyline::Editor;
+use tokio::signal;
 use url::Url;
 
 /// run and execute SQL statements and commands, against a context with the given print options
 pub async fn exec_from_commands(
     ctx: &mut SessionContext,
-    print_options: &PrintOptions,
     commands: Vec<String>,
+    print_options: &PrintOptions,
 ) {
     for sql in commands {
         match exec_and_print(ctx, print_options, sql).await {
@@ -104,8 +105,8 @@ pub async fn exec_from_lines(
 }
 
 pub async fn exec_from_files(
-    files: Vec<String>,
     ctx: &mut SessionContext,
+    files: Vec<String>,
     print_options: &PrintOptions,
 ) {
     let files = files
@@ -165,9 +166,15 @@ pub async fn exec_from_repl(
             }
             Ok(line) => {
                 rl.add_history_entry(line.trim_end())?;
-                match exec_and_print(ctx, print_options, line).await {
-                    Ok(_) => {}
-                    Err(err) => eprintln!("{err}"),
+                tokio::select! {
+                    res = exec_and_print(ctx, print_options, line) => match res {
+                        Ok(_) => {}
+                        Err(err) => eprintln!("{err}"),
+                    },
+                    _ = signal::ctrl_c() => {
+                        println!("^C");
+                        continue
+                    },
                 }
                 // dialect might have changed
                 rl.helper_mut().unwrap().set_dialect(

datafusion-cli/src/main.rs

@@ -216,7 +216,7 @@ pub async fn main() -> Result<()> {
 
     if commands.is_empty() && files.is_empty() {
         if !rc.is_empty() {
-            exec::exec_from_files(rc, &mut ctx, &print_options).await
+            exec::exec_from_files(&mut ctx, rc, &print_options).await
         }
         // TODO maybe we can have thiserror for cli but for now let's keep it simple
         return exec::exec_from_repl(&mut ctx, &mut print_options)
@@ -225,11 +225,11 @@ pub async fn main() -> Result<()> {
     }
 
     if !files.is_empty() {
-        exec::exec_from_files(files, &mut ctx, &print_options).await;
+        exec::exec_from_files(&mut ctx, files, &print_options).await;
     }
 
     if !commands.is_empty() {
-        exec::exec_from_commands(&mut ctx, &print_options, commands).await;
+        exec::exec_from_commands(&mut ctx, commands, &print_options).await;
     }
 
     Ok(())

datafusion-cli/src/print_options.rs

@@ -141,7 +141,8 @@ impl PrintOptions {
         let mut row_count = 0_usize;
         let mut with_header = true;
 
-        while let Some(Ok(batch)) = stream.next().await {
+        while let Some(maybe_batch) = stream.next().await {
+            let batch = maybe_batch?;
             row_count += batch.num_rows();
             self.format.print_batches(
                 &mut writer,

datafusion-examples/README.md

@@ -62,7 +62,9 @@ cargo run --example csv_sql
 - [`simple_udf.rs`](examples/simple_udf.rs): Define and invoke a User Defined Scalar Function (UDF)
 - [`advanced_udf.rs`](examples/advanced_udf.rs): Define and invoke a more complicated User Defined Scalar Function (UDF)
 - [`simple_udaf.rs`](examples/simple_udaf.rs): Define and invoke a User Defined Aggregate Function (UDAF)
+- [`advanced_udaf.rs`](examples/advanced_udaf.rs): Define and invoke a more complicated User Defined Aggregate Function (UDAF)
 - [`simple_udfw.rs`](examples/simple_udwf.rs): Define and invoke a User Defined Window Function (UDWF)
+- [`advanced_udwf.rs`](examples/advanced_udwf.rs): Define and invoke a more complicated User Defined Window Function (UDWF)
 
 ## Distributed