blob: 60baf9bd61a628e16ce2380ddac96327b680fec8 [file]
// 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.
//! # MATCH_RECOGNIZE Example
//!
//! This example demonstrates implementing SQL `MATCH_RECOGNIZE` pattern matching
//! using a custom [`RelationPlanner`]. Unlike the [`pivot_unpivot`] example that
//! rewrites SQL to standard operations, this example creates a **custom logical
//! plan node** (`MiniMatchRecognizeNode`) to represent the operation.
//!
//! ## Supported Syntax
//!
//! ```sql
//! SELECT * FROM events
//! MATCH_RECOGNIZE (
//! PARTITION BY region
//! MEASURES SUM(price) AS total, AVG(price) AS average
//! PATTERN (A B+ C)
//! DEFINE
//! A AS price < 100,
//! B AS price BETWEEN 100 AND 200,
//! C AS price > 200
//! ) AS matches
//! ```
//!
//! ## Architecture
//!
//! This example demonstrates **logical planning only**. Physical execution would
//! require implementing an [`ExecutionPlan`] (see the [`table_sample`] example
//! for a complete implementation with physical planning).
//!
//! ```text
//! SQL Query
//! │
//! ▼
//! ┌─────────────────────────────────────┐
//! │ MatchRecognizePlanner │
//! │ (RelationPlanner trait) │
//! │ │
//! │ • Parses MATCH_RECOGNIZE syntax │
//! │ • Creates MiniMatchRecognizeNode │
//! │ • Converts SQL exprs to DataFusion │
//! └─────────────────────────────────────┘
//! │
//! ▼
//! ┌─────────────────────────────────────┐
//! │ MiniMatchRecognizeNode │
//! │ (UserDefinedLogicalNode) │
//! │ │
//! │ • measures: [(alias, expr), ...] │
//! │ • definitions: [(symbol, expr), ...]│
//! └─────────────────────────────────────┘
//! ```
//!
//! [`pivot_unpivot`]: super::pivot_unpivot
//! [`table_sample`]: super::table_sample
//! [`ExecutionPlan`]: datafusion::physical_plan::ExecutionPlan
use std::{any::Any, cmp::Ordering, hash::Hasher, sync::Arc};
use arrow::array::{ArrayRef, Float64Array, Int32Array, StringArray};
use arrow::record_batch::RecordBatch;
use datafusion::prelude::*;
use datafusion_common::{DFSchemaRef, Result};
use datafusion_expr::{
Expr, UserDefinedLogicalNode,
logical_plan::{Extension, InvariantLevel, LogicalPlan},
planner::{
PlannedRelation, RelationPlanner, RelationPlannerContext, RelationPlanning,
},
};
use datafusion_sql::sqlparser::ast::TableFactor;
use insta::assert_snapshot;
// ============================================================================
// Example Entry Point
// ============================================================================
/// Runs the MATCH_RECOGNIZE examples demonstrating pattern matching on event streams.
///
/// Note: This example demonstrates **logical planning only**. Physical execution
/// would require additional implementation of an [`ExecutionPlan`].
pub async fn match_recognize() -> Result<()> {
let ctx = SessionContext::new();
ctx.register_relation_planner(Arc::new(MatchRecognizePlanner))?;
register_sample_data(&ctx)?;
println!("MATCH_RECOGNIZE Example (Logical Planning Only)");
println!("================================================\n");
run_examples(&ctx).await
}
async fn run_examples(ctx: &SessionContext) -> Result<()> {
// Example 1: Basic MATCH_RECOGNIZE with MEASURES and DEFINE
// Demonstrates: Aggregate measures over matched rows
let plan = run_example(
ctx,
"Example 1: MATCH_RECOGNIZE with aggregations",
r#"SELECT * FROM events
MATCH_RECOGNIZE (
PARTITION BY 1
MEASURES SUM(price) AS total_price, AVG(price) AS avg_price
PATTERN (A)
DEFINE A AS price > 10
) AS matches"#,
)
.await?;
assert_snapshot!(plan, @r"
Projection: matches.price
SubqueryAlias: matches
MiniMatchRecognize measures=[total_price := sum(events.price), avg_price := avg(events.price)] define=[a := events.price > Int64(10)]
TableScan: events
");
// Example 2: Stock price pattern detection
// Demonstrates: Real-world use case finding prices above threshold
let plan = run_example(
ctx,
"Example 2: Detect high stock prices",
r#"SELECT * FROM stock_prices
MATCH_RECOGNIZE (
MEASURES
MIN(price) AS min_price,
MAX(price) AS max_price,
AVG(price) AS avg_price
PATTERN (HIGH)
DEFINE HIGH AS price > 151.0
) AS trends"#,
)
.await?;
assert_snapshot!(plan, @r"
Projection: trends.symbol, trends.price
SubqueryAlias: trends
MiniMatchRecognize measures=[min_price := min(stock_prices.price), max_price := max(stock_prices.price), avg_price := avg(stock_prices.price)] define=[high := stock_prices.price > Float64(151)]
TableScan: stock_prices
");
Ok(())
}
/// Helper to run a single example query and display the logical plan.
async fn run_example(ctx: &SessionContext, title: &str, sql: &str) -> Result<String> {
println!("{title}:\n{sql}\n");
let plan = ctx.sql(sql).await?.into_unoptimized_plan();
let plan_str = plan.display_indent().to_string();
println!("{plan_str}\n");
Ok(plan_str)
}
/// Register test data tables.
fn register_sample_data(ctx: &SessionContext) -> Result<()> {
// events: simple price series
ctx.register_batch(
"events",
RecordBatch::try_from_iter(vec![(
"price",
Arc::new(Int32Array::from(vec![5, 12, 8, 15, 20])) as ArrayRef,
)])?,
)?;
// stock_prices: realistic stock data
ctx.register_batch(
"stock_prices",
RecordBatch::try_from_iter(vec![
(
"symbol",
Arc::new(StringArray::from(vec!["DDOG", "DDOG", "DDOG", "DDOG"]))
as ArrayRef,
),
(
"price",
Arc::new(Float64Array::from(vec![150.0, 155.0, 152.0, 158.0])),
),
])?,
)?;
Ok(())
}
// ============================================================================
// Logical Plan Node: MiniMatchRecognizeNode
// ============================================================================
/// A custom logical plan node representing MATCH_RECOGNIZE operations.
///
/// This is a simplified implementation that captures the essential structure:
/// - `measures`: Aggregate expressions computed over matched rows
/// - `definitions`: Symbol definitions (predicate expressions)
///
/// A production implementation would also include:
/// - Pattern specification (regex-like pattern)
/// - Partition and order by clauses
/// - Output mode (ONE ROW PER MATCH, ALL ROWS PER MATCH)
/// - After match skip strategy
#[derive(Debug)]
struct MiniMatchRecognizeNode {
input: Arc<LogicalPlan>,
schema: DFSchemaRef,
/// Measures: (alias, aggregate_expr)
measures: Vec<(String, Expr)>,
/// Symbol definitions: (symbol_name, predicate_expr)
definitions: Vec<(String, Expr)>,
}
impl UserDefinedLogicalNode for MiniMatchRecognizeNode {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"MiniMatchRecognize"
}
fn inputs(&self) -> Vec<&LogicalPlan> {
vec![&self.input]
}
fn schema(&self) -> &DFSchemaRef {
&self.schema
}
fn check_invariants(&self, _check: InvariantLevel) -> Result<()> {
Ok(())
}
fn expressions(&self) -> Vec<Expr> {
self.measures
.iter()
.chain(&self.definitions)
.map(|(_, expr)| expr.clone())
.collect()
}
fn fmt_for_explain(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "MiniMatchRecognize")?;
if !self.measures.is_empty() {
write!(f, " measures=[")?;
for (i, (alias, expr)) in self.measures.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{alias} := {expr}")?;
}
write!(f, "]")?;
}
if !self.definitions.is_empty() {
write!(f, " define=[")?;
for (i, (symbol, expr)) in self.definitions.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{symbol} := {expr}")?;
}
write!(f, "]")?;
}
Ok(())
}
fn with_exprs_and_inputs(
&self,
exprs: Vec<Expr>,
inputs: Vec<LogicalPlan>,
) -> Result<Arc<dyn UserDefinedLogicalNode>> {
let expected_len = self.measures.len() + self.definitions.len();
if exprs.len() != expected_len {
return Err(datafusion_common::plan_datafusion_err!(
"MiniMatchRecognize: expected {expected_len} expressions, got {}",
exprs.len()
));
}
let input = inputs.into_iter().next().ok_or_else(|| {
datafusion_common::plan_datafusion_err!(
"MiniMatchRecognize requires exactly one input"
)
})?;
let (measure_exprs, definition_exprs) = exprs.split_at(self.measures.len());
let measures = self
.measures
.iter()
.zip(measure_exprs)
.map(|((alias, _), expr)| (alias.clone(), expr.clone()))
.collect();
let definitions = self
.definitions
.iter()
.zip(definition_exprs)
.map(|((symbol, _), expr)| (symbol.clone(), expr.clone()))
.collect();
Ok(Arc::new(Self {
input: Arc::new(input),
schema: Arc::clone(&self.schema),
measures,
definitions,
}))
}
fn dyn_hash(&self, state: &mut dyn Hasher) {
state.write_usize(Arc::as_ptr(&self.input) as usize);
state.write_usize(self.measures.len());
state.write_usize(self.definitions.len());
}
fn dyn_eq(&self, other: &dyn UserDefinedLogicalNode) -> bool {
other.as_any().downcast_ref::<Self>().is_some_and(|o| {
Arc::ptr_eq(&self.input, &o.input)
&& self.measures == o.measures
&& self.definitions == o.definitions
})
}
fn dyn_ord(&self, other: &dyn UserDefinedLogicalNode) -> Option<Ordering> {
if self.dyn_eq(other) {
Some(Ordering::Equal)
} else {
None
}
}
}
// ============================================================================
// Relation Planner: MatchRecognizePlanner
// ============================================================================
/// Relation planner that creates `MiniMatchRecognizeNode` for MATCH_RECOGNIZE queries.
#[derive(Debug)]
struct MatchRecognizePlanner;
impl RelationPlanner for MatchRecognizePlanner {
fn plan_relation(
&self,
relation: TableFactor,
ctx: &mut dyn RelationPlannerContext,
) -> Result<RelationPlanning> {
let TableFactor::MatchRecognize {
table,
measures,
symbols,
alias,
..
} = relation
else {
return Ok(RelationPlanning::Original(relation));
};
// Plan the input table
let input = ctx.plan(*table)?;
let schema = input.schema().clone();
// Convert MEASURES: SQL expressions → DataFusion expressions
let planned_measures: Vec<(String, Expr)> = measures
.iter()
.map(|m| {
let alias = ctx.normalize_ident(m.alias.clone());
let expr = ctx.sql_to_expr(m.expr.clone(), schema.as_ref())?;
Ok((alias, expr))
})
.collect::<Result<_>>()?;
// Convert DEFINE: symbol definitions → DataFusion expressions
let planned_definitions: Vec<(String, Expr)> = symbols
.iter()
.map(|s| {
let name = ctx.normalize_ident(s.symbol.clone());
let expr = ctx.sql_to_expr(s.definition.clone(), schema.as_ref())?;
Ok((name, expr))
})
.collect::<Result<_>>()?;
// Create the custom node
let node = MiniMatchRecognizeNode {
input: Arc::new(input),
schema,
measures: planned_measures,
definitions: planned_definitions,
};
let plan = LogicalPlan::Extension(Extension {
node: Arc::new(node),
});
Ok(RelationPlanning::Planned(PlannedRelation::new(plan, alias)))
}
}