Merge branch 'main' into perf-octet-length

Author: Brijesh-Thakkar

.github/workflows/audit.yml

@@ -42,7 +42,7 @@ jobs:
     steps:
       - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8  # v6.0.1
       - name: Install cargo-audit
-        uses: taiki-e/install-action@e0db384ad69f5ba2c6dd0129d8934e0d0ba465c1  # v2.65.10
+        uses: taiki-e/install-action@a983ca795126a4be4e8a44879ac5c4c3ef66cae1  # v2.65.11
         with:
           tool: cargo-audit
       - name: Run audit check

.github/workflows/rust.yml

@@ -421,7 +421,7 @@ jobs:
           sudo apt-get update -qq
           sudo apt-get install -y -qq clang
       - name: Setup wasm-pack
-        uses: taiki-e/install-action@e0db384ad69f5ba2c6dd0129d8934e0d0ba465c1  # v2.65.10
+        uses: taiki-e/install-action@a983ca795126a4be4e8a44879ac5c4c3ef66cae1  # v2.65.11
         with:
           tool: wasm-pack
       - name: Run tests with headless mode
@@ -741,7 +741,7 @@ jobs:
       - name: Setup Rust toolchain
         uses: ./.github/actions/setup-builder
       - name: Install cargo-msrv
-        uses: taiki-e/install-action@e0db384ad69f5ba2c6dd0129d8934e0d0ba465c1  # v2.65.10
+        uses: taiki-e/install-action@a983ca795126a4be4e8a44879ac5c4c3ef66cae1  # v2.65.11
         with:
           tool: cargo-msrv
 

Cargo.lock

@@ -185,8 +185,7 @@ sqlparser = { version = "0.59.0", default-features = false, features = ["std", "
 strum = "0.27.2"
 strum_macros = "0.27.2"
 tempfile = "3"
-testcontainers = { version = "0.25.2", features = ["default"] }
-testcontainers-modules = { version = "0.13" }
+testcontainers-modules = { version = "0.14" }
 tokio = { version = "1.48", features = ["macros", "rt", "sync"] }
 url = "2.5.7"
 zstd = { version = "0.13", default-features = false }

Cargo.toml

@@ -74,5 +74,4 @@ ctor = { workspace = true }
 insta = { workspace = true }
 insta-cmd = "0.6.0"
 rstest = { workspace = true }
-testcontainers = { workspace = true }
 testcontainers-modules = { workspace = true, features = ["minio"] }

datafusion-cli/Cargo.toml

@@ -24,10 +24,12 @@ use insta::{Settings, glob};
 use insta_cmd::{assert_cmd_snapshot, get_cargo_bin};
 use std::path::PathBuf;
 use std::{env, fs};
-use testcontainers::core::{CmdWaitFor, ExecCommand, Mount};
-use testcontainers::runners::AsyncRunner;
-use testcontainers::{ContainerAsync, ImageExt, TestcontainersError};
 use testcontainers_modules::minio;
+use testcontainers_modules::testcontainers::core::{CmdWaitFor, ExecCommand, Mount};
+use testcontainers_modules::testcontainers::runners::AsyncRunner;
+use testcontainers_modules::testcontainers::{
+    ContainerAsync, ImageExt, TestcontainersError,
+};
 
 fn cli() -> Command {
     Command::new(get_cargo_bin("datafusion-cli"))

datafusion-cli/tests/cli_integration.rs

@@ -83,13 +83,12 @@ use std::{
     any::Any,
     fmt::{self, Debug, Formatter},
     hash::{Hash, Hasher},
-    ops::{Add, Div, Mul, Sub},
     pin::Pin,
-    str::FromStr,
     sync::Arc,
     task::{Context, Poll},
 };
 
+use arrow::datatypes::{Float64Type, Int64Type};
 use arrow::{
     array::{ArrayRef, Int32Array, RecordBatch, StringArray, UInt32Array},
     compute,
@@ -102,6 +101,7 @@ use futures::{
 use rand::{Rng, SeedableRng, rngs::StdRng};
 use tonic::async_trait;
 
+use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal;
 use datafusion::{
     execution::{
         RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder,
@@ -410,11 +410,12 @@ impl RelationPlanner for TableSamplePlanner {
                 "TABLESAMPLE requires a quantity (percentage, fraction, or row count)"
             );
         };
+        let quantity_value_expr = context.sql_to_expr(quantity.value, input.schema())?;
 
         match quantity.unit {
             // TABLESAMPLE (N ROWS) - exact row limit
             Some(TableSampleUnit::Rows) => {
-                let rows = parse_quantity::<i64>(&quantity.value)?;
+                let rows: i64 = parse_literal::<Int64Type>(&quantity_value_expr)?;
                 if rows < 0 {
                     return plan_err!("row count must be non-negative, got {}", rows);
                 }
@@ -426,15 +427,15 @@ impl RelationPlanner for TableSamplePlanner {
 
             // TABLESAMPLE (N PERCENT) - percentage sampling
             Some(TableSampleUnit::Percent) => {
-                let percent = parse_quantity::<f64>(&quantity.value)?;
+                let percent: f64 = parse_literal::<Float64Type>(&quantity_value_expr)?;
                 let fraction = percent / 100.0;
                 let plan = TableSamplePlanNode::new(input, fraction, seed).into_plan();
                 Ok(RelationPlanning::Planned(PlannedRelation::new(plan, alias)))
             }
 
             // TABLESAMPLE (N) - fraction if <1.0, row limit if >=1.0
             None => {
-                let value = parse_quantity::<f64>(&quantity.value)?;
+                let value = parse_literal::<Float64Type>(&quantity_value_expr)?;
                 if value < 0.0 {
                     return plan_err!("sample value must be non-negative, got {}", value);
                 }
@@ -453,40 +454,6 @@ impl RelationPlanner for TableSamplePlanner {
     }
 }
 
-/// Parse a SQL expression as a numeric value (supports basic arithmetic).
-fn parse_quantity<T>(expr: &ast::Expr) -> Result<T>
-where
-    T: FromStr + Add<Output = T> + Sub<Output = T> + Mul<Output = T> + Div<Output = T>,
-{
-    eval_numeric_expr(expr)
-        .ok_or_else(|| plan_datafusion_err!("invalid numeric expression: {:?}", expr))
-}
-
-/// Recursively evaluate numeric SQL expressions.
-fn eval_numeric_expr<T>(expr: &ast::Expr) -> Option<T>
-where
-    T: FromStr + Add<Output = T> + Sub<Output = T> + Mul<Output = T> + Div<Output = T>,
-{
-    match expr {
-        ast::Expr::Value(v) => match &v.value {
-            ast::Value::Number(n, _) => n.to_string().parse().ok(),
-            _ => None,
-        },
-        ast::Expr::BinaryOp { left, op, right } => {
-            let l = eval_numeric_expr::<T>(left)?;
-            let r = eval_numeric_expr::<T>(right)?;
-            match op {
-                ast::BinaryOperator::Plus => Some(l + r),
-                ast::BinaryOperator::Minus => Some(l - r),
-                ast::BinaryOperator::Multiply => Some(l * r),
-                ast::BinaryOperator::Divide => Some(l / r),
-                _ => None,
-            }
-        }
-        _ => None,
-    }
-}
-
 /// Custom logical plan node representing a TABLESAMPLE operation.
 ///
 /// Stores sampling parameters (bounds, seed) and wraps the input plan.

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -107,6 +107,8 @@ impl DataFrameWriteOptions {
     }
 
     /// Set the single_file_output value to true or false
+    ///
+    /// When set to true, an output file will always be created even if the DataFrame is empty
     pub fn with_single_file_output(mut self, single_file_output: bool) -> Self {
         self.single_file_output = single_file_output;
         self

datafusion/core/src/dataframe/mod.rs

@@ -17,3 +17,96 @@
 
 //! Re-exports the [`datafusion_datasource_arrow::file_format`] module, and contains tests for it.
 pub use datafusion_datasource_arrow::file_format::*;
+
+#[cfg(test)]
+mod tests {
+    use futures::StreamExt;
+    use std::sync::Arc;
+
+    use arrow::array::{Int64Array, StringArray};
+    use arrow::datatypes::{DataType, Field, Schema};
+    use arrow::record_batch::RecordBatch;
+    use datafusion_common::Result;
+
+    use crate::execution::options::ArrowReadOptions;
+    use crate::prelude::SessionContext;
+
+    #[tokio::test]
+    async fn test_write_empty_arrow_from_sql() -> Result<()> {
+        let ctx = SessionContext::new();
+
+        let tmp_dir = tempfile::TempDir::new()?;
+        let path = format!("{}/empty_sql.arrow", tmp_dir.path().to_string_lossy());
+
+        ctx.sql(&format!(
+            "COPY (SELECT CAST(1 AS BIGINT) AS id LIMIT 0) TO '{path}' STORED AS ARROW",
+        ))
+        .await?
+        .collect()
+        .await?;
+
+        assert!(std::path::Path::new(&path).exists());
+
+        let read_df = ctx.read_arrow(&path, ArrowReadOptions::default()).await?;
+        let stream = read_df.execute_stream().await?;
+
+        assert_eq!(stream.schema().fields().len(), 1);
+        assert_eq!(stream.schema().field(0).name(), "id");
+
+        let results: Vec<_> = stream.collect().await;
+        let total_rows: usize = results
+            .iter()
+            .filter_map(|r| r.as_ref().ok())
+            .map(|b| b.num_rows())
+            .sum();
+        assert_eq!(total_rows, 0);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_write_empty_arrow_from_record_batch() -> Result<()> {
+        let ctx = SessionContext::new();
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int64, false),
+            Field::new("name", DataType::Utf8, true),
+        ]));
+        let empty_batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(Int64Array::from(Vec::<i64>::new())),
+                Arc::new(StringArray::from(Vec::<Option<&str>>::new())),
+            ],
+        )?;
+
+        let tmp_dir = tempfile::TempDir::new()?;
+        let path = format!("{}/empty_batch.arrow", tmp_dir.path().to_string_lossy());
+
+        ctx.register_batch("empty_table", empty_batch)?;
+
+        ctx.sql(&format!("COPY empty_table TO '{path}' STORED AS ARROW"))
+            .await?
+            .collect()
+            .await?;
+
+        assert!(std::path::Path::new(&path).exists());
+
+        let read_df = ctx.read_arrow(&path, ArrowReadOptions::default()).await?;
+        let stream = read_df.execute_stream().await?;
+
+        assert_eq!(stream.schema().fields().len(), 2);
+        assert_eq!(stream.schema().field(0).name(), "id");
+        assert_eq!(stream.schema().field(1).name(), "name");
+
+        let results: Vec<_> = stream.collect().await;
+        let total_rows: usize = results
+            .iter()
+            .filter_map(|r| r.as_ref().ok())
+            .map(|b| b.num_rows())
+            .sum();
+        assert_eq!(total_rows, 0);
+
+        Ok(())
+    }
+}