Expose parser on DFParser to enable user controlled parsing

[tshauck]

Which issue does this PR close?

Closes #533

Rationale for this change

Being able to control the underlying parser is useful to be able to add capabilities between the SQL and logical plan layer, which can propagate downstream in the plan.

This approach seems to work nicely insofar as modifying how the SQL becomes a user-defined Statement. I.e. here I'm wrapping DF's Statement with my own, and creating it based on the underlying result from parse_copy.

However, I'm unsure on the next steps. E.g. with this strategy of wrapping the statement, I can't then use DF's machinery to generate the LogicalPlan, like using sql_statement_to_plan. It's not too bad in the basic case, as it'd be easy to generate a user defined logical node from this enum (or fall back to sql_statement_to_plan), but in more complex cases there's probably a better way, e.g. maybe to make Statement a Trait with a visit method, or something along those lines? There's probably a better/simpler approach :).

Going to open this in draft, and will come back later today/tomorrow to see what the basic strategy for constructing a LogicalPlan would look like and see if that's a reasonable first chunk.

Open to feedback on any of it, thanks!

What changes are included in this PR?

Makes the parser on DFParser public and adds an example.

Are these changes tested?

Manually ran the example. Seems to perform as expected. Though not sure what the best path is next.

Are there any user-facing changes?

Not breaking, but parser is available on DFParser.

[tshauck]

@alamb, would you be able to please give me some feedback here?

In a7a7808, I made parser public and added an example, and that seemed to go well insofar as I could create a new MyStatement which is an enum of a DF Statement and a custom one.

That was good, but then it can't work with df functionality like sql_statement_to_plan, so in 86588e4 I started down added an Extension enum variant to Statement, based on how the LogicalPlan extensions work, but I'm having trouble integrating it into DF writ large... e.g. I'm running into many object safety issues...

So I guess my questions are 1) does an Extension statement seem reasonable to you? 2) and if so, any advice on the implementation to mitigate object safety issues.

[alamb]

Thanks for working on this @tshauck -- it is really cool to see

Going to open this in draft, and will come back later today/tomorrow to see what the basic strategy for constructing a LogicalPlan would look like and see if that's a reasonable first chunk.

The simplest strategy I can think of is to implement your own version to sql_statement_to_plan but for MyStatement

So something like

    pub async fn my_statement_to_plan(
        ctx: &SessionContext,
        statement: MyStatement,
    ) -> Result<LogicalPlan> {
   ...
   // call statement_to_plan here and implement the special Copy logic
}

So I guess my questions are 1) does an Extension statement seem reasonable to you? 2) and if so, any advice on the implementation to mitigate object safety issues.

I don't fully understand the need for an Extension statement -- if the idea is to wrap DFStatement with MyStatement I think you can just implement the planning in terms of MyStatement as above.

If the idea is to avoid repetition with whatever statement_to_plan is doing, maybe we can factor out the common functionality into a module.

[tshauck]

Thanks for the feedback, I reverted the extension changes, and now it only has the example plus making Parser public on DFParser. To your point later, there's quite a bit within sql_statement_to_plan (and functions it calls) that would be nice to not have to replicate, but perhaps going through the work of using this change in my code will make it clearer what those changes should be. Things like resolving name, handling subqueries, etc.

I guess in terms of next steps, I could move this out of draft or add my_statement_to_plan to the example, if you think that's a more complete change?

[alamb]

Thank you @tshauck -- this looks great 🙏 . It is really nice to get an example showing how to do this.

I left some suggestions on additional contextual comments and how to make the code potentially simpler, but I also think we could merge this PR as is.

Can you also add an entry to the README for this example?
https://github.com/apache/arrow-datafusion/tree/main/datafusion-examples#single-process

Also, I think we can make the example code simpler with less nesting with a function. Something like

impl MyParser<'_> {
...
    /// Returns true if the next token is `COPY` keyword, false otherwise
    fn is_copy(&self) -> bool {
        matches!(
            self.df_parser.parser.peek_token().token,
            Token::Word(w) if w.keyword == Keyword::COPY
        )
    }

    pub fn parse_statement(&mut self) -> Result<MyStatement, ParserError> {
        if self.is_copy() {
            self.df_parser.parser.next_token(); // COPY
            let df_statement = self.df_parser.parse_copy()?;

            if let Statement::CopyTo(s) = df_statement {
                Ok(MyStatement::from(s))
            } else {
                Ok(MyStatement::DFStatement(Box::from(df_statement)))
            }
        } else {
            let df_statement = self.df_parser.parse_statement()?;
            Ok(MyStatement::from(df_statement))
        }
    }
...

[alamb]

😸

[tshauck]

Here's sql_dialect.rs -- I also alphabetized the list, hope that's ok.

[tshauck]

@alamb Thanks for all the feedback! The matches! approach is really nice... re-requesting review.

[alamb]

[alamb]

❤️ -- very nice drive by cleanup 🙏

[alamb]

Thanks again @tshauck