Skip to content

Bug: calling "with_new_exprs" on join after optimization unexpectedly fails #14999

@niebayes

Description

@niebayes

Describe the bug

Before optimization, specifically the ExtractEquijoinPredicate rule, calling with_new_exprs succeeds.
However, after optimized by ExtractEquijoinPredicate, calling with_new_exprs unexpectedly fails with the following error:

            "The front part expressions should be an binary equality expression, actual:t1.a"

To Reproduce

 #[tokio::test]
    async fn test_join_with_new_exprs() -> Result<()> {
        // Creates a default session context.
        let session_context = SessionContext::new();

        // Registers two tables.
        session_context.register_table(
            "t1",
            Arc::new(EmptyTable::new(
                Schema::new(vec![Field::new("a", DataType::Int32, true)]).into(),
            )),
        )?;
        session_context.register_table(
            "t2",
            Arc::new(EmptyTable::new(
                Schema::new(vec![Field::new("a", DataType::Int32, true)]).into(),
            )),
        )?;

        // Creates a plan consisting of a join operator and extracts the join operator.
        let join = session_context
            .sql("select * from t1 join t2 on t1.a = t2.a")
            .await?
            .logical_plan()
            .inputs()
            .remove(0)
            .clone();
        assert!(matches!(join, LogicalPlan::Join(_)));

        let LogicalPlan::Join(before) = &join else {
            unreachable!()
        };
        // Should be: "[]".
        println!("on before opt: {:?}", before.on);
        // Should be: "Some(BinaryExpr(BinaryExpr { left: Column(Column { relation: Some(Bare { table: "t1" }), name: "a" }), op: Eq,
        //                                          right: Column(Column { relation: Some(Bare { table: "t2" }), name: "a" }) }))".
        println!("filter before opt: {:?}", before.filter);
        // Invokes `with_new_exprs` before optimization should succeed.
        let res = join.with_new_exprs(
            join.expressions(),
            join.inputs().into_iter().map(|x| x.clone()).collect(),
        );
        assert!(res.is_ok());

        print!("\n\n");

        // Optimizes join.
        let join = session_context.state().optimize(&join)?;
        let LogicalPlan::Join(after) = &join else {
            unreachable!()
        };
        // Should be: "[(Column(Column { relation: Some(Bare { table: "t1" }), name: "a" }), Column(Column { relation: Some(Bare { table: "t2" }), name: "a" }))]".
        println!("on after opt: {:?}", after.on);
        // Should be: "None".
        println!("filter after opt: {:?}", after.filter);
        // Invokes `with_new_exprs` after optimization unexpectedly fails.
        let res = join.with_new_exprs(
            join.expressions(),
            join.inputs().into_iter().map(|x| x.clone()).collect(),
        );
        assert!(res.is_err());
        assert_contains!(
            res.unwrap_err().to_string(),
            "The front part expressions should be an binary equality expression, actual:t1.a"
        );

        Ok(())
    }

Expected behavior

No response

Additional context

No response

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions