Add relation planner extension support to customize SQL planning

datafusion-examples/Cargo.toml

@@ -46,10 +46,14 @@ dashmap = { workspace = true }
 # note only use main datafusion crate for examples
 base64 = "0.22.1"
 datafusion = { workspace = true, default-features = true, features = ["parquet_encryption"] }
+datafusion-common = { workspace = true }
+datafusion-expr = { workspace = true }
 datafusion-physical-expr-adapter = { workspace = true }
 datafusion-proto = { workspace = true }
+datafusion-sql = { workspace = true }
 env_logger = { workspace = true }
 futures = { workspace = true }
+insta = { workspace = true }
 log = { workspace = true }
 mimalloc = { version = "0.1", default-features = false }
 object_store = { workspace = true, features = ["aws", "http"] }

datafusion-examples/README.md

@@ -86,6 +86,9 @@ cargo run --example dataframe -- dataframe
 - [`examples/external_dependency/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
+- [`examples/relation_planner/match_recognize.rs`](examples/relation_planner/match_recognize.rs): Use custom relation planner to implement MATCH_RECOGNIZE pattern matching
+- [`examples/relation_planner/pivot_unpivot.rs`](examples/relation_planner/pivot_unpivot.rs): Use custom relation planner to implement PIVOT and UNPIVOT operations
+- [`examples/relation_planner/table_sample.rs`](examples/relation_planner/table_sample.rs): Use custom relation planner to implement TABLESAMPLE clause
 - [`examples/data_io/remote_catalog.rs`](examples/data_io/remote_catalog.rs): Examples of interfacing with a remote catalog (e.g. over a network)
 - [`examples/udf/simple_udaf.rs`](examples/udf/simple_udaf.rs): Define and invoke a User Defined Aggregate Function (UDAF)
 - [`examples/udf/simple_udf.rs`](examples/udf/simple_udf.rs): Define and invoke a User Defined Scalar Function (UDF)

datafusion-examples/examples/relation_planner/main.rs

@@ -0,0 +1,141 @@
+// 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.
+
+//! # Relation Planner Examples
+//!
+//! These examples demonstrate how to use custom relation planners to extend
+//! DataFusion's SQL syntax with custom table operators.
+//!
+//! ## Usage
+//! ```bash
+//! cargo run --example relation_planner -- [match_recognize|pivot_unpivot|table_sample]
+//! ```
+//!
+//! Each subcommand runs a corresponding example:
+//! - `match_recognize` — MATCH_RECOGNIZE pattern matching on event streams
+//! - `pivot_unpivot` — PIVOT and UNPIVOT operations for reshaping data
+//! - `table_sample` — TABLESAMPLE clause for sampling rows from tables
+//!
+//! ## Snapshot Testing
+//!
+//! These examples use [insta](https://insta.rs) for inline snapshot assertions.
+//! If query output changes, regenerate the snapshots with:
+//! ```bash
+//! cargo insta test --example relation_planner --accept
+//! ```
+
+mod match_recognize;
+mod pivot_unpivot;
+mod table_sample;
+
+use std::str::FromStr;
+
+use datafusion::error::{DataFusionError, Result};
+
+enum ExampleKind {
+    MatchRecognize,
+    PivotUnpivot,
+    TableSample,
+}
+
+impl AsRef<str> for ExampleKind {
+    fn as_ref(&self) -> &str {
+        match self {
+            Self::MatchRecognize => "match_recognize",
+            Self::PivotUnpivot => "pivot_unpivot",
+            Self::TableSample => "table_sample",
+        }
+    }
+}
+
+impl FromStr for ExampleKind {
+    type Err = DataFusionError;
+
+    fn from_str(s: &str) -> Result<Self> {
+        match s {
+            "match_recognize" => Ok(Self::MatchRecognize),
+            "pivot_unpivot" => Ok(Self::PivotUnpivot),
+            "table_sample" => Ok(Self::TableSample),
+            _ => Err(DataFusionError::Execution(format!("Unknown example: {s}"))),
+        }
+    }
+}
+
+impl ExampleKind {
+    const ALL: [Self; 3] = [Self::MatchRecognize, Self::PivotUnpivot, Self::TableSample];
+
+    const EXAMPLE_NAME: &str = "relation_planner";
+
+    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())
+    })?;
+
+    if arg == "all" {
+        for example in ExampleKind::ALL {
+            match example {
+                ExampleKind::MatchRecognize => match_recognize::match_recognize().await?,
+                ExampleKind::PivotUnpivot => pivot_unpivot::pivot_unpivot().await?,
+                ExampleKind::TableSample => table_sample::table_sample().await?,
+            }
+        }
+    } else {
+        match arg.parse::<ExampleKind>()? {
+            ExampleKind::MatchRecognize => match_recognize::match_recognize().await?,
+            ExampleKind::PivotUnpivot => pivot_unpivot::pivot_unpivot().await?,
+            ExampleKind::TableSample => table_sample::table_sample().await?,
+        }
+    }
+
+    Ok(())
+}
+
+/// Test wrappers that enable `cargo insta test --example relation_planner --accept`
+/// to regenerate inline snapshots. Without these, insta cannot run the examples
+/// in test mode since they only have `main()` functions.
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_match_recognize() {
+        match_recognize::match_recognize().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_pivot_unpivot() {
+        pivot_unpivot::pivot_unpivot().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_table_sample() {
+        table_sample::table_sample().await.unwrap();
+    }
+}