Make PruningPredicate's rewrite public

[adriangb]

Replaces #12606, #12835

As per #12606 (comment) the plan was to split the rewrite logic from the rest of PruningPredicate. It turns out that was as easy as making a function public, kudos to whoever wrote this originally for organizing it nicely 🥳.

[adriangb]

cc @alamb

[alamb]

Thanks @adriangb -- this looks like a great idea to me -- I left a comment on how we can potentially improve the API slightly. Let me know what you think

[alamb]

I think you can write literals more concisely like this if you wanted

Suggested change

	let input = col("b").eq(lit(ScalarValue::Int32(Some(12))));
	let input = col("b").eq(lit(12i32));

[alamb]

Since this is a new API I think it might be worth thinking a bit more about. What would you think about moving this logic into its own structure?

For example, instead of

        let known_expression = col("a").eq(lit(ScalarValue::Int32(Some(12))));
        let known_expression_transformed = rewrite_predicate_to_statistics_predicate(
            &logical2physical(&known_expression, &schema),
            &schema,
            None,
        );

Something like

        let known_expression = col("a").eq(lit(ScalarValue::Int32(Some(12))));
        let known_expression_transformed = StatisticsPredicateRewriter::new()
           .rewrite(
             &logical2physical(&known_expression, &schema)
             &schema)
           );

Or with a rewrite hook

            let expr = logical2physical(&expr, &schema_with_b);
            StatisticsPredicateRewriter::new()
              .with_unhandledhook(|expr| {
                 // closure that handles expr
              })
           .rewrite(
                &expr,
                &schema,
            )

[adriangb]

Yup that sounds good to me, I'll push that change tomorrow.

[adriangb]

Actually just did it, it was easy.

My main question is if it should include the schema or not.

[alamb]

Thanks @adriangb -- this looks good to me.

My main question is if it should include the schema or not.

I think callers could probably get the schema in other ways. However, it might be nice to pass it in too.

[adriangb]

I think callers could probably get the schema in other ways. However, it might be nice to pass it in too.

I just meant:

        let known_expression_transformed = StatisticsPredicateRewriter::new()
           .rewrite(
             &logical2physical(&known_expression, &schema)
             &schema)
           );

vs

        let known_expression_transformed = StatisticsPredicateRewriter::new(schema)
           .rewrite(&logical2physical(&known_expression, &schema));

I think both work, I don't see a reason to change it but was just double checking.

[alamb]

Here are some small suggestions: pydantic#3

[adriangb]

Thanks Andrew, merged your suggestions 😄

[alamb]

Doh, looks like I broke clippy

https://github.com/apache/datafusion/actions/runs/11294824353/job/31416484653?pr=12850

505 | impl ConstantUnhandledPredicateHook {
    | ----------------------------------- associated function in this implementation
506 |     fn new(default: Arc<dyn PhysicalExpr>) -> Self {
    |        ^^^

[adriangb]

I can fix it or feel free to just push to this branch directly.

[alamb]

I can fix it or feel free to just push to this branch directly.

I don't think I can as it is the pydanitic fork and I don't have access to do so

[adriangb]

Ah right I think you would have been able to if it was my personal fork but you're right it's the pydantic one.

I think I fixed it now.

[alamb]

Thanks again @adriangb