application of simple optimizer rule produces incorrect results (DF 49 regression)

[wkalt]

Describe the bug

The following code produces an incorrect result on Datafusion 49.0.2, but correct behavior on 48.0.1. Both execution plans should produce the same row count, but in DF 49.0.2 one returns too few results. The behavior is independent of the order in which the original/optimized plans are executed. If the order of the executions is flipped, the second execution always returns the incorrect result.

use arrow_array::{RecordBatch, UInt32Array};
use arrow_schema::{DataType, Field, Schema};
use datafusion::datasource::memory::MemTable;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::*;
use datafusion_common::Result as DFResult;
use std::sync::Arc;

#[derive(Debug)]
struct NoOpOptimizer;

impl PhysicalOptimizerRule for NoOpOptimizer {
    fn optimize(
        &self,
        plan: Arc<dyn ExecutionPlan>,
        _config: &datafusion_common::config::ConfigOptions,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        Ok(plan)
    }

    fn name(&self) -> &str {
        "NoOpOptimizer"
    }

    fn schema_check(&self) -> bool {
        true
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create simple test data
    let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)]));

    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![Arc::new(UInt32Array::from((0..20).collect::<Vec<u32>>()))],
    )?;

    let ctx = SessionContext::new();
    let provider = MemTable::try_new(schema.clone(), vec![vec![batch]])?;
    ctx.register_table("test", Arc::new(provider))?;

    // Create a simple plan with a sort and a limit.
    let df = ctx
        .table("test")
        .await?
        .sort(vec![col("id").sort(true, true)])?
        .limit(0, Some(10))?;

    let logical_plan = df.into_optimized_plan()?;
    let plan = ctx.state().create_physical_plan(&logical_plan).await?;

    // Apply no-op optimizer
    let optimizer = NoOpOptimizer;
    let optimized_plan = optimizer.optimize(plan.clone(), &Default::default())?;

    let task_ctx = ctx.task_ctx();

    let original_stream = plan.execute(0, task_ctx.clone())?;
    let original_batches = datafusion::physical_plan::common::collect(original_stream).await?;
    let original_count: usize = original_batches.iter().map(|b| b.num_rows()).sum();

    let optimized_stream = optimized_plan.execute(0, task_ctx)?;
    let optimized_batches = datafusion::physical_plan::common::collect(optimized_stream).await?;
    let optimized_count: usize = optimized_batches.iter().map(|b| b.num_rows()).sum();

    // Result mismatch on DF 49
    println!("Expected: 10 rows (limit in query)");
    println!("Optimized plan: {} rows", optimized_count);
    println!("Original plan:  {} rows", original_count);

    Ok(())
}

Cargo.toml:

[package]
name = "datafusion-regression-test"
version = "0.1.0"
edition = "2021"

[workspace]

[dependencies]
arrow-array = "55"
arrow-schema = "55"
datafusion = "49.0.2"
datafusion-common = "49.0.2"
tokio = { version = "1", features = ["full"] }

DF 49.0.2:

Expected: 10 rows (limit in query)
Optimized plan: 9 rows
Original plan:  10 rows

DF 48.01:

Expected: 10 rows (limit in query)
Optimized plan: 10 rows
Original plan:  10 rows

To Reproduce

No response

Expected behavior

No response

Additional context

No response

[wkalt]

reading more commits, I suspect this may be #16998 which is fixed by #17028. If so the fix is just to update our code to use reset_state. In that case this can be closed.

[wkalt]

Calling reset_state between executions did not seem to address the issue, so I'll leave it open for comment.

Our usecase for doing this is testing optimizer rules, to make sure the result is the same before and after optimization. I have worked around this by changing the test to create two separate plans.

[alamb]

@adriangb any ideas of what is going on here?

[adriangb]

I don't think this has anything to do with dynamic filters.

I reduced the MRE to:

#[tokio::test]
async fn test_regression() {
    use arrow::{
        array::{RecordBatch, UInt32Array},
        datatypes::{DataType, Field, Schema},
    };
    use datafusion::datasource::memory::MemTable;
    use datafusion::prelude::*;
    use std::sync::Arc;

    // Create simple test data
    let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)]));

    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![Arc::new(UInt32Array::from((0..20).collect::<Vec<u32>>()))],
    ).unwrap();

    let cfg = SessionConfig::new().set_bool("datafusion.optimizer.enable_dynamic_filter_pushdown", false);
    let ctx = SessionContext::new_with_config(cfg);
    let provider = MemTable::try_new(schema.clone(), vec![vec![batch]]).unwrap();
    ctx.register_table("test", Arc::new(provider)).unwrap();

    // Create a simple plan with a sort and a limit.
    let df = ctx
        .table("test")
        .await.unwrap()
        .sort(vec![col("id").sort(true, true)]).unwrap()
        .limit(0, Some(10)).unwrap();

    let logical_plan = df.into_optimized_plan().unwrap();
    let plan_1 = ctx.state().create_physical_plan(&logical_plan).await.unwrap();
    let plan_2 = plan_1.clone();
    let task_ctx = ctx.task_ctx();

    let stream_1 = plan_1.execute(0, task_ctx.clone()).unwrap();
    let batches_1 = datafusion::physical_plan::common::collect(stream_1).await.unwrap();
    let count_1: usize = batches_1.iter().map(|b| b.num_rows()).sum();

    // let plan_2 = plan_2.reset_state().unwrap();

    let stream_2 = plan_2.execute(0, task_ctx).unwrap();
    let batches_2 = datafusion::physical_plan::common::collect(stream_2).await.unwrap();
    let count_2: usize = batches_2.iter().map(|b| b.num_rows()).sum();

    // Result mismatch on DF 49
    println!("Expected: 10 rows (limit in query)");
    assert_eq!(count_2, count_1, "Row count mismatch: expected {count_1}, got {count_2}");
}

(I tested by adding this to the bottom of datafusion/core/tests/physical_optimizer/filter_pushdown/mod.rs)

If you comment back in the line // let plan_2 = plan_2.reset_state().unwrap(); then it passes.

As you can see I set datafusion.optimizer.enable_dynamic_filter_pushdown = false which should disable dynamic filters, i.e. this is unrelated to that feature.

Could it be that you are just observing the state of the SortExec being re-used across runs? Where did you call reset_state()?

[alamb]

when I add this code to the original reproducer the test now passes.

    let optimized_plan = optimized_plan.reset_state().unwrap();

Using @adriangb 's reproducer, when I uncommented out the line

    // let plan_2 = plan_2.reset_state().unwrap();

The test passes for me on DataFusion 49 and when I comment it back out then the test fails

Thus, I think the solution is to call reset_state.

The root cause seems to be that code doesn't really copy the plan so this example is running the same plan twice.

This line specifically:

    let optimized_plan = optimizer.optimize(plan.clone(), &Default::default())?;

plan is an Arc<dyn ExecutionPlan> so calling clone on it just copies the pointer

Similarly, in @adriangb 's version plan_1 is an Arc<dyn ExecutionPlan> so calling clone on it just copies the pointer

    let plan_1 = ctx.state().create_physical_plan(&logical_plan).await.unwrap();
    let plan_2 = plan_1.clone();

So in effect what the reproducers are doing is calling execute() on the same plan twice. Thus I would expect that calling reset_state on the plan would be needed to reset any internal state, which is exactly the behavior we see

[alamb]

I think we can close this ticket as not a bug, is that ok @wkalt ?

[wkalt]

@alamb sorry for the delay, yes that is fine. Closing.

[alamb]

Thank you for the report and the reproducer -- I am glad we sorted this out