Merge branch 'main' into spark_abs

Author: hsiang-c

.github/workflows/audit.yml

@@ -42,7 +42,7 @@ jobs:
     steps:
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
       - name: Install cargo-audit
-        uses: taiki-e/install-action@6cc14f7f2f4b3129aff07a8b071d2d4f2733465d  # v2.62.50
+        uses: taiki-e/install-action@0be4756f42223b67aa4b7df5effad59010cbf4b9  # v2.62.51
         with:
           tool: cargo-audit
       - name: Run audit check

.github/workflows/rust.yml

@@ -271,7 +271,21 @@ jobs:
     runs-on: ubuntu-latest
     container:
       image: amd64/rust
+      volumes:
+        - /usr/local:/host/usr/local
     steps:
+      - name: Remove unnecessary preinstalled software
+        run: |
+          echo "Disk space before cleanup:"
+          df -h
+          # remove tool cache: about 8.5GB (github has host /opt/hostedtoolcache mounted as /__t)
+          rm -rf /__t/* || true
+          # remove Haskell runtime: about 6.3GB (host /usr/local/.ghcup)
+          rm -rf /host/usr/local/.ghcup || true
+          # remove Android library: about 7.8GB (host /usr/local/lib/android)
+          rm -rf /host/usr/local/lib/android || true
+          echo "Disk space after cleanup:"
+          df -h
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
         with:
           submodules: true
@@ -434,7 +448,7 @@ jobs:
           sudo apt-get update -qq
           sudo apt-get install -y -qq clang
       - name: Setup wasm-pack
-        uses: taiki-e/install-action@6cc14f7f2f4b3129aff07a8b071d2d4f2733465d  # v2.62.50
+        uses: taiki-e/install-action@0be4756f42223b67aa4b7df5effad59010cbf4b9  # v2.62.51
         with:
           tool: wasm-pack
       - name: Run tests with headless mode
@@ -761,7 +775,7 @@ jobs:
       - name: Setup Rust toolchain
         uses: ./.github/actions/setup-builder
       - name: Install cargo-msrv
-        uses: taiki-e/install-action@6cc14f7f2f4b3129aff07a8b071d2d4f2733465d  # v2.62.50
+        uses: taiki-e/install-action@0be4756f42223b67aa4b7df5effad59010cbf4b9  # v2.62.51
         with:
           tool: cargo-msrv
 
@@ -806,4 +820,4 @@ jobs:
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
         with:
           persist-credentials: false
-      - uses: crate-ci/typos@1af53e3774f068183ffd0c7193eb061a2b65a531  # v1.39.1
+      - uses: crate-ci/typos@626c4bedb751ce0b7f03262ca97ddda9a076ae1c  # v1.39.2

datafusion-examples/README.md

@@ -51,7 +51,7 @@ cargo run --example dataframe
 - [`examples/udf/advanced_udwf.rs`](examples/udf/advanced_udwf.rs): Define and invoke a more complicated User Defined Window Function (UDWF)
 - [`examples/data_io/parquet_advanced_index.rs`](examples/data_io/parquet_advanced_index.rs): Creates a detailed secondary index that covers the contents of several parquet files
 - [`examples/udf/async_udf.rs`](examples/udf/async_udf.rs): Define and invoke an asynchronous User Defined Scalar Function (UDF)
-- [`analyzer_rule.rs`](examples/analyzer_rule.rs): Use a custom AnalyzerRule to change a query's semantics (row level access control)
+- [`examples/query_planning/analyzer_rule.rs`](examples/query_planning/analyzer_rule.rs): Use a custom AnalyzerRule to change a query's semantics (row level access control)
 - [`examples/data_io/catalog.rs`](examples/data_io/catalog.rs): Register the table into a custom catalog
 - [`examples/data_io/json_shredding.rs`](examples/data_io/json_shredding.rs): Shows how to implement custom filter rewriting for JSON shredding
 - [`composed_extension_codec`](examples/composed_extension_codec.rs): Example of using multiple extension codecs for serialization / deserialization
@@ -65,22 +65,23 @@ cargo run --example dataframe
 - [`examples/builtin_functions/date_time`](examples/builtin_functions/date_time.rs): Examples of date-time related functions and queries
 - [`default_column_values.rs`](examples/default_column_values.rs): Implement custom default value handling for missing columns using field metadata and PhysicalExprAdapter
 - [`deserialize_to_struct.rs`](examples/deserialize_to_struct.rs): Convert query results (Arrow ArrayRefs) into Rust structs
-- [`expr_api.rs`](examples/expr_api.rs): Create, execute, simplify, analyze and coerce `Expr`s
+- [`examples/query_planning/expr_api.rs`](examples/query_planning/expr_api.rs): Create, execute, simplify, analyze and coerce `Expr`s
 - [`examples/custom_data_source/file_stream_provider.rs`](examples/custom_data_source/file_stream_provider.rs): Run a query on `FileStreamProvider` which implements `StreamProvider` for reading and writing to arbitrary stream sources / sinks.
 - [`flight/sql_server.rs`](examples/flight/sql_server.rs): Run DataFusion as a standalone process and execute SQL queries from Flight and and FlightSQL (e.g. JDBC) clients
 - [`examples/builtin_functions/function_factory.rs`](examples/builtin_functions/function_factory.rs): Register `CREATE FUNCTION` handler to implement SQL macros
 - [`memory_pool_tracking.rs`](examples/memory_pool_tracking.rs): Demonstrates TrackConsumersPool for memory tracking and debugging with enhanced error messages
 - [`memory_pool_execution_plan.rs`](examples/memory_pool_execution_plan.rs): Shows how to implement memory-aware ExecutionPlan with memory reservation and spilling
-- [`optimizer_rule.rs`](examples/optimizer_rule.rs): Use a custom OptimizerRule to replace certain predicates
+- [`examples/query_planning/optimizer_rule.rs`](examples/query_planning/optimizer_rule.rs): Use a custom OptimizerRule to replace certain predicates
 - [`examples/data_io/parquet_embedded_index.rs`](examples/data_io/parquet_embedded_index.rs): Store a custom index inside a Parquet file and use it to speed up queries
 - [`examples/data_io/parquet_encrypted.rs`](examples/data_io/parquet_encrypted.rs): Read and write encrypted Parquet files using DataFusion
 - [`examples/data_io/parquet_encrypted_with_kms.rs`](examples/data_io/parquet_encrypted_with_kms.rs): Read and write encrypted Parquet files using an encryption factory
 - [`examples/data_io/parquet_index.rs`](examples/data_io/parquet_index.rs): Create an secondary index over several parquet files and use it to speed up queries
 - [`examples/data_io/parquet_exec_visitor.rs`](examples/data_io/parquet_exec_visitor.rs): Extract statistics by visiting an ExecutionPlan after execution
-- [`parse_sql_expr.rs`](examples/parse_sql_expr.rs): Parse SQL text into DataFusion `Expr`.
-- [`plan_to_sql.rs`](examples/plan_to_sql.rs): Generate SQL from DataFusion `Expr` and `LogicalPlan`
-- [`planner_api.rs`](examples/planner_api.rs) APIs to manipulate logical and physical plans
-- [`pruning.rs`](examples/pruning.rs): Use pruning to rule out files based on statistics
+- [`examples/query_planning/parse_sql_expr.rs`](examples/query_planning/parse_sql_expr.rs): Parse SQL text into DataFusion `Expr`.
+- [`examples/query_planning/plan_to_sql.rs`](examples/query_planning/plan_to_sql.rs): Generate SQL from DataFusion `Expr` and `LogicalPlan`
+- [`examples/query_planning/planner_api.rs`](examples/query_planning/planner_api.rs) APIs to manipulate logical and physical plans
+- [`examples/query_planning/pruning.rs`](examples/query_planning/pruning.rs): Use pruning to rule out files based on statistics
+- [`examples/query_planning/thread_pools.rs`](examples/query_planning/thread_pools.rs): Demonstrates TrackConsumersPool for memory tracking and debugging with enhanced error messages and shows how to implement memory-aware ExecutionPlan with memory reservation and spilling
 - [`query-aws-s3.rs`](examples/external_dependency/query-aws-s3.rs): Configure `object_store` and run a query against files stored in AWS S3
 - [`examples/data_io/query_http_csv.rs`](examples/data_io/query_http_csv.rs): Configure `object_store` and run a query against files via HTTP
 - [`examples/builtin_functions/regexp.rs`](examples/builtin_functions/regexp.rs): Examples of using regular expression functions

datafusion-examples/examples/analyzer_rule.rs renamed to datafusion-examples/examples/query_planning/analyzer_rule.rs

@@ -35,8 +35,7 @@ use std::sync::{Arc, Mutex};
 /// level access control scheme by introducing a filter to the query.
 ///
 /// See [optimizer_rule.rs] for an example of a optimizer rule
-#[tokio::main]
-pub async fn main() -> Result<()> {
+pub async fn analyzer_rule() -> Result<()> {
     // AnalyzerRules run before OptimizerRules.
     //
     // DataFusion includes several built in AnalyzerRules for tasks such as type

datafusion-examples/examples/expr_api.rs renamed to datafusion-examples/examples/query_planning/expr_api.rs

@@ -55,8 +55,7 @@ use datafusion::prelude::*;
 /// 5. Analyze predicates for boundary ranges: [`range_analysis_demo`]
 /// 6. Get the types of the expressions: [`expression_type_demo`]
 /// 7. Apply type coercion to expressions: [`type_coercion_demo`]
-#[tokio::main]
-async fn main() -> Result<()> {
+pub async fn expr_api() -> Result<()> {
     // The easiest way to do create expressions is to use the
     // "fluent"-style API:
     let expr = col("a") + lit(5);

datafusion-examples/examples/query_planning/main.rs

@@ -0,0 +1,138 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! # These are all internal mechanics of the query planning and optimization layers
+//!
+//! These examples demonstrate internal mechanics of the query planning and optimization layers.
+//!
+//! ## Usage
+//! ```bash
+//! cargo run --example query_planning -- [analyzer_rule|expr_api|optimizer_rule|parse_sql_expr|plan_to_sql|planner_api|pruning|thread_pools]
+//! ```
+//!
+//! Each subcommand runs a corresponding example:
+//! - `analyzer_rule` — use a custom AnalyzerRule to change a query's semantics (row level access control)
+//! - `expr_api` — create, execute, simplify, analyze and coerce `Expr`s
+//! - `optimizer_rule` — use a custom OptimizerRule to replace certain predicates
+//! - `parse_sql_expr` — parse SQL text into DataFusion `Expr`
+//! - `plan_to_sql` — generate SQL from DataFusion `Expr` and `LogicalPlan`
+//! - `planner_api` — APIs to manipulate logical and physical plans
+//! - `pruning` — APIs to manipulate logical and physical plans
+//! - `thread_pools` — demonstrate TrackConsumersPool for memory tracking and debugging with enhanced error messages and shows how to implement memory-aware ExecutionPlan with memory reservation and spilling
+
+mod analyzer_rule;
+mod expr_api;
+mod optimizer_rule;
+mod parse_sql_expr;
+mod plan_to_sql;
+mod planner_api;
+mod pruning;
+mod thread_pools;
+
+use std::str::FromStr;
+
+use datafusion::error::{DataFusionError, Result};
+
+enum ExampleKind {
+    AnalyzerRule,
+    ExprApi,
+    OptimizerRule,
+    ParseSqlExpr,
+    PlanToSql,
+    PlannerApi,
+    Pruning,
+    ThreadPools,
+}
+
+impl AsRef<str> for ExampleKind {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::AnalyzerRule => "analyzer_rule",
+            Self::ExprApi => "expr_api",
+            Self::OptimizerRule => "optimizer_rule",
+            Self::ParseSqlExpr => "parse_sql_expr",
+            Self::PlanToSql => "plan_to_sql",
+            Self::PlannerApi => "planner_api",
+            Self::Pruning => "pruning",
+            Self::ThreadPools => "thread_pools",
+        }
+    }
+}
+
+impl FromStr for ExampleKind {
+    type Err = DataFusionError;
+
+    fn from_str(s: &str) -> Result<Self> {
+        match s {
+            "analyzer_rule" => Ok(Self::AnalyzerRule),
+            "expr_api" => Ok(Self::ExprApi),
+            "optimizer_rule" => Ok(Self::OptimizerRule),
+            "parse_sql_expr" => Ok(Self::ParseSqlExpr),
+            "plan_to_sql" => Ok(Self::PlanToSql),
+            "planner_api" => Ok(Self::PlannerApi),
+            "pruning" => Ok(Self::Pruning),
+            "thread_pools" => Ok(Self::ThreadPools),
+            _ => Err(DataFusionError::Execution(format!("Unknown example: {s}"))),
+        }
+    }
+}
+
+impl ExampleKind {
+    const ALL: [Self; 8] = [
+        Self::AnalyzerRule,
+        Self::ExprApi,
+        Self::OptimizerRule,
+        Self::ParseSqlExpr,
+        Self::PlanToSql,
+        Self::PlannerApi,
+        Self::Pruning,
+        Self::ThreadPools,
+    ];
+
+    const EXAMPLE_NAME: &str = "query_planning";
+
+    fn variants() -> Vec<&'static str> {
+        Self::ALL.iter().map(|x| x.as_ref()).collect()
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let usage = format!(
+        "Usage: cargo run --example {} -- [{}]",
+        ExampleKind::EXAMPLE_NAME,
+        ExampleKind::variants().join("|")
+    );
+
+    let arg = std::env::args().nth(1).ok_or_else(|| {
+        eprintln!("{usage}");
+        DataFusionError::Execution("Missing argument".to_string())
+    })?;
+
+    match arg.parse::<ExampleKind>()? {
+        ExampleKind::AnalyzerRule => analyzer_rule::analyzer_rule().await?,
+        ExampleKind::ExprApi => expr_api::expr_api().await?,
+        ExampleKind::OptimizerRule => optimizer_rule::optimizer_rule().await?,
+        ExampleKind::ParseSqlExpr => parse_sql_expr::parse_sql_expr().await?,
+        ExampleKind::PlanToSql => plan_to_sql::plan_to_sql_examples().await?,
+        ExampleKind::PlannerApi => planner_api::planner_api().await?,
+        ExampleKind::Pruning => pruning::pruning().await?,
+        ExampleKind::ThreadPools => thread_pools::thread_pools().await?,
+    }
+
+    Ok(())
+}

datafusion-examples/examples/optimizer_rule.rs renamed to datafusion-examples/examples/query_planning/optimizer_rule.rs

@@ -37,8 +37,7 @@ use std::sync::Arc;
 ///
 /// See [analyzer_rule.rs] for an example of AnalyzerRules, which are for
 /// changing plan semantics.
-#[tokio::main]
-pub async fn main() -> Result<()> {
+pub async fn optimizer_rule() -> Result<()> {
     // DataFusion includes many built in OptimizerRules for tasks such as outer
     // to inner join conversion and constant folding.
     //

datafusion-examples/examples/parse_sql_expr.rs renamed to datafusion-examples/examples/query_planning/parse_sql_expr.rs

@@ -32,17 +32,15 @@ use datafusion::{
 /// The code in this example shows how to:
 ///
 /// 1. [`simple_session_context_parse_sql_expr_demo`]: Parse a simple SQL text into a logical
-/// expression using a schema at [`SessionContext`].
+///    expression using a schema at [`SessionContext`].
 ///
 /// 2. [`simple_dataframe_parse_sql_expr_demo`]: Parse a simple SQL text into a logical expression
-/// using a schema at [`DataFrame`].
+///    using a schema at [`DataFrame`].
 ///
 /// 3. [`query_parquet_demo`]: Query a parquet file using the parsed_sql_expr from a DataFrame.
 ///
 /// 4. [`round_trip_parse_sql_expr_demo`]: Parse a SQL text and convert it back to SQL using [`Unparser`].
-
-#[tokio::main]
-async fn main() -> Result<()> {
+pub async fn parse_sql_expr() -> Result<()> {
     // See how to evaluate expressions
     simple_session_context_parse_sql_expr_demo()?;
     simple_dataframe_parse_sql_expr_demo().await?;

datafusion-examples/examples/plan_to_sql.rs renamed to datafusion-examples/examples/query_planning/plan_to_sql.rs

@@ -43,28 +43,26 @@ use std::sync::Arc;
 /// The code in this example shows how to:
 ///
 /// 1. [`simple_expr_to_sql_demo`]: Create a simple expression [`Exprs`] with
-/// fluent API and convert to sql suitable for passing to another database
+///    fluent API and convert to sql suitable for passing to another database
 ///
 /// 2. [`simple_expr_to_pretty_sql_demo`] Create a simple expression
-/// [`Exprs`] with fluent API and convert to sql without extra parentheses,
-/// suitable for displaying to humans
+///    [`Exprs`] with fluent API and convert to sql without extra parentheses,
+///    suitable for displaying to humans
 ///
 /// 3. [`simple_expr_to_sql_demo_escape_mysql_style`]" Create a simple
-/// expression [`Exprs`] with fluent API and convert to sql escaping column
-/// names in MySQL style.
+///    expression [`Exprs`] with fluent API and convert to sql escaping column
+///    names in MySQL style.
 ///
 /// 4. [`simple_plan_to_sql_demo`]: Create a simple logical plan using the
-/// DataFrames API and convert to sql string.
+///    DataFrames API and convert to sql string.
 ///
 /// 5. [`round_trip_plan_to_sql_demo`]: Create a logical plan from a SQL string, modify it using the
-/// DataFrames API and convert it back to a  sql string.
+///    DataFrames API and convert it back to a  sql string.
 ///
 /// 6. [`unparse_my_logical_plan_as_statement`]: Create a custom logical plan and unparse it as a statement.
 ///
 /// 7. [`unparse_my_logical_plan_as_subquery`]: Create a custom logical plan and unparse it as a subquery.
-
-#[tokio::main]
-async fn main() -> Result<()> {
+pub async fn plan_to_sql_examples() -> Result<()> {
     // See how to evaluate expressions
     simple_expr_to_sql_demo()?;
     simple_expr_to_pretty_sql_demo()?;

datafusion-examples/examples/planner_api.rs renamed to datafusion-examples/examples/query_planning/planner_api.rs

@@ -32,8 +32,7 @@ use datafusion::prelude::*;
 /// physical plan:
 /// - Via the combined `create_physical_plan` API.
 /// - Utilizing the analyzer, optimizer, and query planner APIs separately.
-#[tokio::main]
-async fn main() -> Result<()> {
+pub async fn planner_api() -> Result<()> {
     // Set up a DataFusion context and load a Parquet file
     let ctx = SessionContext::new();
     let testdata = datafusion::test_util::parquet_test_data();