Initial support for prepared statement

[NGA-TRAN]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.
In order to support Prepare statement I am thinking about adding 2 new functions. This is is preliminary thinking. If you have some idea please advice.

Function 1: Compile a SQL with question marks/parameters

• Input: a SQL with question marks
• Output:
  • Logical plan that understands question marks/parameter (not sure how yet)
  • List of data type of the parameters of the question marks

Example:

• Input:

  SELECT  a, b FROM t WHERE x > ? and Y = ?;

• Output:
  • Logical plan that understands question marks/parameter
  • [ i64, String ] where i64 and String are data types of X and Y respectively

Function 2: Run the compiled SQL with provided values of the parameters

• Input:
  • Logical plan that understands question marks/parameter
  • List of values of the corresponding question marks
• Output
  • Physical plan of the logical plan with all question marks replaced with the corresponding values in the list

Same example above

• Input
  • Logical plan that understands question marks/parameter generated for SQL above
  • [ 7, "Boston" ]
• Output:
  • Phychical plan of SELECT a, b FROM t WHERE x > 7 and Y = "Boston"

Describe the solution you'd like
TBD

Describe alternatives you've considered
TBD

Additional context

[NGA-TRAN]

@alamb What do you think?

[alamb]

This is great @NGA-TRAN -- thank you for writing it up. My only feedback is that for option 2 it might be easier if the output was a new LogicalPlan that had the parameter values filled in, which can then be fed through the rest of the planning machinery.

[NGA-TRAN]

Thanks @alamb . Let me see how Logical Plan looks like and propose a clearer one

[alamb]

I think the feature described in this proposal is needed to properly handle prepared statements in FlightSQL

For example, in ballista parameter handling appears to still be a "TODO":
https://github.com/apache/arrow-ballista/blob/acbdf608807697e72122a249c4bf7505274906d0/ballista/scheduler/src/flight_sql.rs#L876

As more background, we are in the process of designing and implementing FlightSQL in IOx and proper parameterized query support is something we have identified as important

cc @avantgardnerio @yahoNanJing @liukun4515 @thinkharderdev @metesynnada / @retikulum -- I wonder if you have any thoughts or ideas about how useful this would be for Ballista or if we have missed some prior art / implementation somewhere

[thinkharderdev]

This is great @NGA-TRAN -- thank you for writing it up. My only feedback is that for option 2 it might be easier if the output was a new LogicalPlan that had the parameter values filled in, which can then be fed through the rest of the planning machinery.

I agree. Conceptually it seems like the bind variable can just be an Expr but physical planning should fail if it hits any bind variable exprs. So before physical planning you rewrite the plan to replace any such expressions with user-supplied literals.

[NGA-TRAN]

Thanks @alamb and @thinkharderdev for the feedback. Here is my proposal.

We will use Postgres's prepared statement syntax

Let us create a table for our examples

CREATE TABLE cpu ("user" INT, region VARCHAR, time TMESTAMP);

-- insert some data

SELECT * FROM cpu;
 user | region |        time         
------+--------+---------------------
    1 | east   | 2022-11-15 00:00:00
    2 | west   | 2022-11-30 19:10:25
(2 rows)

New Feature: PREPARE a select statement and EXECUTE it

We will support PREPARE and EXECUTE statement that are compatible with Postgres as follows:

Example 1:

-- Create a prepared plan
PREPARE my_plan_1(TIMESTAMP) AS 
   SELECT "user", region FROM cpu WHERE time = $1;

-- Execute the plan
EXECUTE my_plan_1('2022-11-30 19:10:25');
 user | region 
------+--------
    2 | west
(1 row)

Example 2

-- Create a prepared plan
PREPARE my_plan_2(INT, INT) AS 
  SELECT region, count(*) FROM cpu GROUP BY region HAVING count(*) >= $1 and count(*) < $2 order by 1;

-- Execute the plan
EXECUTE my_plan_2(1, 10);
 region | count 
--------+-------
 east   |     1
 west   |     1
(2 rows)

Design proposal

Prepare Statement

sqlparser parses and provides output for PREPARE as follows

   Prepare {
        name: Ident,
        data_types: Vec<DataType>,
        statement: Box<Statement>,
    },

We will generate the corresponding struct when we run that statement:

PreparedPlan {
   name: Indent,                            // name provided with the PREPARE
   params: Vec<DataType>,         // list of data types provided with the PREPARE
   plan:   LogicalPlan,                    // logical plan that understands Params defined below
}

struct Param {
   type: DataType,
   idx:   usize             // index to the vector `params`
}

Before showing examples of PreparedPlan, let us look into today logical plans.

-- Example 1:
SELECT “user”, region FROM cpu WHERE time > to_timestamp('1970-01-01T00:00:00.000000120+00:00');

Logical Plan:
 Projection: cpu.user, cpu.region
  Filter: cpu.time > totimestamp(Utf8("1970-01-01T00:00:00.000000120+00:00"))
    TableScan: cpu

-- Example 2:
SELECT region, count(*) FROM cpu GROUP BY region HAVING count(*) >= 1 and count(*) < 10 order by 1;

logical plan:
 Sort: cpu.region ASC NULLS LAST
  Projection: cpu.region, COUNT(UInt8(1))
    Filter: COUNT(UInt8(1)) >= Int64(1) AND COUNT(UInt8(1)) < Int64(10)
      Aggregate: groupBy=[[cpu.region]], aggr=[[COUNT(UInt8(1))]]
        TableScan: cpu

Let see the full output of a prepared statement:

Example 1:

PREPARE my_plan_1(TIMESTAMP) AS 
   SELECT "user", region FROM cpu WHERE time = $1;

The output will look like

my_plan_1 = PreparedPlan {
    name: "my_plan_1",
    params: [Timestamp]     // vector of 1
    plan: Logical Plan
                  Projection: cpu.user, cpu.region
                     Filter: cpu.time > Param { type: Timestamp, idx: 0}   // idx points to the first item of the params
                        TableScan: cpu
}

Example 2:

PREPARE my_plan_2(INT, INT) AS 
  SELECT region, count(*) FROM cpu GROUP BY region HAVING count(*) >= $1 and count(*) < $2 order by 1;

The output will look like

my_plan_2 = PreparedPlan {
   name: "my_plan_2",
   params: [Int64, Int64],
   plan: Logical plan 
               Sort: cpu.region ASC NULLS LAST
                  Projection: cpu.region, COUNT(UInt8(1))
                     Filter: COUNT(UInt8(1)) >= Param { type: Int64, idx: 0 }    // points to first item of the params
                                AND COUNT(UInt8(1)) < Param { type: Int64, idx: 1 }  // points to second items of the params
                         Aggregate: groupBy=[[cpu.region]], aggr=[[COUNT(UInt8(1))]]
                           TableScan: cpu
}

Question: How do we store all the generated prepared plans? In a hashMap?

Execute Plan

sqlparser parses and provides output for EXECUTE as follows

  Execute { name: Ident, parameters: Vec<Expr> },

We will search the the plan and replace Params in the Logical plan with the provided parameters. Then generate physical plan for it and execute.

[NGA-TRAN]

I will start working on this while addressing review comments on design. My check list will be:

• sqlparser parses the PREPARE and EXECUTE (available, just need to make sure they work in DF)
• Generate and store PreparedPlans (many steps that will needs it own checklist)
• Execute a PreparedPlan (this may need a few steps too)

[alamb]

Thanks @NGA-TRAN -- like the proposal in general. I think the trickiest part will be "where to store the prepared statement" (aka my_plan_2)

And I think it is critical for this feature to be opt in (so downstream crates can decide if/how they want to handle the state management)

I suggest following the model of CreateExternalTable

https://github.com/apache/arrow-datafusion/blob/23b4495516a5a6246e1732bca4f82bb065b021be/datafusion/expr/src/logical_plan/plan.rs#L1335-L1358

Where the core of DataFusion knows how to create the LogicalPlan::CreateExternalTable but leaves it to the downstream system (aka datafusion-cli or IOx) to actually execute them by connecting them together.

@Dandandan and @andygrove any thoughts?

[NGA-TRAN]

I have tried to parse PREPARE and EXECUTE statements. They both work well

Here is the test and output of `PREPARE`. sqlparser has enum Value::Placeholder to store the parameter. I will use that in LogicalPlan

 fn create_prepared_statement() -> Result<(), ParserError> {
        // positive case
        let sql = "PREPARE my_plan(TIME, INT) AS SELECT region FROM cpu WHERE time = $1 and usage_user > $2";
        let statements = DFParser::parse_sql(sql)?;
        println!("{:#?}", statements[0]);
        assert_eq!(statements.len(), 1);
        Ok(())
    }

Statement(
    Prepare {
        name: Ident {
            value: "my_plan",
            quote_style: None,
        },
        data_types: [
            Time(
                None,
                None,
            ),
            Int(
                None,
            ),
        ],
        statement: Query(
            Query {
                with: None,
                body: Select(
                    Select {
                        distinct: false,
                        top: None,
                        projection: [
                            UnnamedExpr(
                                Identifier(
                                    Ident {
                                        value: "region",
                                        quote_style: None,
                                    },
                                ),
                            ),
                        ],
                        into: None,
                        from: [
                            TableWithJoins {
                                relation: Table {
                                    name: ObjectName(
                                        [
                                            Ident {
                                                value: "cpu",
                                                quote_style: None,
                                            },
                                        ],
                                    ),
                                    alias: None,
                                    args: None,
                                    with_hints: [],
                                },
                                joins: [],
                            },
                        ],
                        lateral_views: [],
                        selection: Some(
                            BinaryOp {
                                left: BinaryOp {
                                    left: Identifier(
                                        Ident {
                                            value: "time",
                                            quote_style: None,
                                        },
                                    ),
                                    op: Eq,
                                    right: Value(
                                        Placeholder(
                                            "$1",
                                        ),
                                    ),
                                },
                                op: And,
                                right: BinaryOp {
                                    left: Identifier(
                                        Ident {
                                            value: "usage_user",
                                            quote_style: None,
                                        },
                                    ),
                                    op: Gt,
                                    right: Value(
                                        Placeholder(
                                            "$2",
                                        ),
                                    ),
                                },
                            },
                        ),
                        group_by: [],
                        cluster_by: [],
                        distribute_by: [],
                        sort_by: [],
                        having: None,
                        qualify: None,
                    },
                ),
                order_by: [],
                limit: None,
                offset: None,
                fetch: None,
                lock: None,
            },
        ),
    },
)

Here is the test and output of `EXECUTE`

#[test]
    fn execute_statement() -> Result<(), ParserError> {
        // positive case
        let sql = "EXECUTE my_plan(1, '2022-11-30')";
        let statements = DFParser::parse_sql(sql)?;
        println!("{:#?}", statements[0]);
        assert_eq!(statements.len(), 1);
        Ok(())
    }

Statement(
    Execute {
        name: Ident {
            value: "my_plan",
            quote_style: None,
        },
        parameters: [
            Value(
                Number(
                    "1",
                    false,
                ),
            ),
            Value(
                SingleQuotedString(
                    "2022-11-30",
                ),
            ),
        ],
    },
)

And here the statement of the query without binding. It is the same as the one with binding but the the PlaceHolder is replaced wit actual value

``` Statement( Query( Query { with: None, body: Select( Select { distinct: false, top: None, projection: [ UnnamedExpr( Identifier( Ident { value: "region", quote_style: None, }, ), ), ], into: None, from: [ TableWithJoins { relation: Table { name: ObjectName( [ Ident { value: "cpu", quote_style: None, }, ], ), alias: None, args: None, with_hints: [], }, joins: [], }, ], lateral_views: [], selection: Some( BinaryOp { left: BinaryOp { left: Identifier( Ident { value: "time", quote_style: None, }, ), op: Eq, right: Value( Number( "10", false, ), ), }, op: And, right: BinaryOp { left: Identifier( Ident { value: "usage_user", quote_style: None, }, ), op: Gt, right: Value( Number( "20", false, ), ), }, }, ), group_by: [], cluster_by: [], distribute_by: [], sort_by: [], having: None, qualify: None, }, ), order_by: [], limit: None, offset: None, fetch: None, lock: None, }, ), ) ```

[metesynnada]

My first understanding is that

• We can utilize new LogicalPlan enums with name

  • CreatePlan
  • ExecutePlan

  this can be implemented in core Datafusion.

• We need a place to store the intermediate plan. This requires unique keys and handling for keyspace and serialization. So, we need proto updates for this intermediate logical plan. We can utilize ObjectStore.put() in datafusion, or using StateBackendClient in Ballista.

• We need some functionality for intermediate plan to actual plan. This can be handled in core Datafusion, and Ballista scheduler.

• Also, I agree with @alamb to handle ExecutePlan outside of the DefaultPhysicalPlanner.

[alamb]

The current functionality split works like this:

SessionContext::create_logical_plan creates a LogicalPlan for DDL statements like CREATE EXTERNAL TABLE and CREATE VIEW

However, SessionContext::create_physical_plan can not create an execution plan for such statements (it errors if you try)

Instead, SessionContext::sql provides some basic implementation for the DDL statements.

https://github.com/apache/arrow-datafusion/blob/40e44a64f51438ea8eae7112828ef116409a15f0/datafusion/core/src/execution/context.rs#L251-L481

So if DataFusion users want the basic implementations, they can use SessionContext::sql

If they want to control / disallow such statements (as we do in IOx, for example), instead one uses the lower level APIs to create the physical plan directly.

I think this is a reasonable model and suggest we follow the same pattern with prepared statements CreatePlan and ExecutePlan as described by @metesynnada

If someone is feeling ambitious, they could try and consolidate the DDL variants to make this split clearer. Here is one proposal: #3349

[liukun4515]

I think the feature described in this proposal is needed to properly handle prepared statements in FlightSQL

For example, in ballista parameter handling appears to still be a "TODO": https://github.com/apache/arrow-ballista/blob/acbdf608807697e72122a249c4bf7505274906d0/ballista/scheduler/src/flight_sql.rs#L876

As more background, we are in the process of designing and implementing FlightSQL in IOx and proper parameterized query support is something we have identified as important

cc @avantgardnerio @yahoNanJing @liukun4515 @thinkharderdev @metesynnada / @retikulum -- I wonder if you have any thoughts or ideas about how useful this would be for Ballista or if we have missed some prior art / implementation somewhere

Sorry for the late reply, It is a great feature.
It will reduce the compile time significantly.

As you @alamb described in the comments, "where to store the prepared logical plan?"

Maybe the hashmap in the memory is enough for the init stage