Skip to content

Commit

Permalink
WIP demonstrate using get_field with expr_api in 37.1.0
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Apr 22, 2024
1 parent 4f5c52a commit 9560732
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 20 deletions.
26 changes: 16 additions & 10 deletions datafusion-examples/examples/expr_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,21 +24,20 @@ use arrow::record_batch::RecordBatch;
use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use datafusion::common::{DFField, DFSchema};
use datafusion::error::Result;
use datafusion::functions_array::rewrite::ArrayFunctionRewriter;
use datafusion::optimizer::simplify_expressions::ExprSimplifier;
use datafusion::physical_expr::{
analyze, create_physical_expr, AnalysisContext, ExprBoundaries, PhysicalExpr,
};
use datafusion::prelude::*;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::TreeNode;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::{ScalarValue, ToDFSchema};
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::expr::BinaryExpr;
use datafusion_expr::expr_rewriter::FunctionRewrite;
use datafusion_expr::interval_arithmetic::Interval;
use datafusion_expr::simplify::SimplifyContext;
use datafusion_expr::{ColumnarValue, ExprSchemable, Operator};
use datafusion_expr::{
ColumnarValue, ExprSchemable, GetFieldAccess, GetIndexedField, Operator,
};

/// This example demonstrates the DataFusion [`Expr`] API.
///
Expand Down Expand Up @@ -310,12 +309,19 @@ pub fn physical_expr(schema: &Schema, expr: Expr) -> Result<Arc<dyn PhysicalExpr
// apply type coercion here to ensure types match
let expr = simplifier.coerce(expr, df_schema.clone())?;

// Support array functions by rewriting expressions
let rewriter = ArrayFunctionRewriter::new();

// Support Expr::struct by rewriting expressions
let expr = expr
.transform_up(|expr| {
rewriter.rewrite(expr, df_schema.as_ref(), &ConfigOptions::default())
.transform_up(&|expr| {
Ok(match expr {
Expr::GetIndexedField(GetIndexedField {
expr,
field: GetFieldAccess::NamedStructField { name },
}) => {
let name = Expr::Literal(name);
Transformed::yes(get_field(*expr, name))
}
_ => Transformed::no(expr),
})
})?
.data;

Expand Down
12 changes: 2 additions & 10 deletions datafusion/functions-array/src/rewrite.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,16 +30,8 @@ use datafusion_expr::expr_rewriter::FunctionRewrite;
use datafusion_expr::{BinaryExpr, Expr, GetFieldAccess, GetIndexedField, Operator};
use datafusion_functions::expr_fn::get_field;

/// Rewrites expressions such as `expr[field]` into function dor array functions
#[derive(Debug, Default)]
pub struct ArrayFunctionRewriter {}

impl ArrayFunctionRewriter {
/// Create a new `ArrayFunctionRewriter`
pub fn new() -> Self {
Self::default()
}
}
/// Rewrites expressions into function calls to array functions
pub(crate) struct ArrayFunctionRewriter {}

impl FunctionRewrite for ArrayFunctionRewriter {
fn name(&self) -> &str {
Expand Down

0 comments on commit 9560732

Please sign in to comment.