column types must match schema types occurs after `unnest_columns` on another column

[ion-elgreco]

Describe the bug

I am rewriting our CDF operation in delta-rs, the code looks roughly like this:

    let mut projected = if should_cdc {
        operation_count
            .clone()
            .with_column(
                CDC_COLUMN_NAME,
                when(col(TARGET_DELETE_COLUMN).is_null(), lit("delete")) // nulls are equal to True
                    .when(col(DELETE_COLUMN).is_null(), lit("source_delete"))
                    .when(col(TARGET_COPY_COLUMN).is_null(), lit("copy"))
                    .when(col(TARGET_INSERT_COLUMN).is_null(), lit("insert"))
                    .when(col(TARGET_UPDATE_COLUMN).is_null(), lit("update"))
                    .end()?,
            )?
            // .drop_columns(&["__delta_rs_path"])? // WEIRD bug caused by interaction with unnest_columns, has to be dropped otherwise throws schema error
            .with_column(
                "__delta_rs_update_expanded",
                when(
                    col(CDC_COLUMN_NAME).eq(lit("update")),
                    lit(ScalarValue::List(ScalarValue::new_list(
                        &[
                            ScalarValue::Utf8(Some("update_preimage".into())),
                            ScalarValue::Utf8(Some("update_postimage".into())),
                        ],
                        &DataType::List(Field::new("element", DataType::Utf8, false).into()),
                        true,
                    ))),
                )
                .end()?,
            )?
            .unnest_columns(&["__delta_rs_update_expanded"])?
            .with_column(
                CDC_COLUMN_NAME,
                when(
                    col(CDC_COLUMN_NAME).eq(lit("update")),
                    col("__delta_rs_update_expanded"),
                )
                .otherwise(col(CDC_COLUMN_NAME))?,
            )?
            .drop_columns(&["__delta_rs_update_expanded"])?
            .select(write_projection_with_cdf)?

I noticed that when I do unnest_columns on another column, it complains afterwards about a schema error:

Result::unwrap()` on an `Err` value: Arrow { source: InvalidArgumentError("column types must match schema types, expected Utf8 but found Dictionary(UInt16, Utf8) at column index 7") }

Since I don't need the column, I can safely drop it beforehand, but I don't understand why doesn't Dictionary(UInt16, Utf8) just coerce to utf8?

To Reproduce

Bit difficult but, you can run grab my branch: https://github.com/ion-elgreco/delta-rs/tree/refactor--combine_execution_plans

And then you run the test test_merge_cdc_enabled_simple, with this line commented out: .drop_columns(&["__delta_rs_path"])?

Expected behavior

I guess coerce gracefully?

Additional context

No response

[duongcongtoai]

take

[duongcongtoai]

duplicate with #6057

[duongcongtoai]

This error happens when the provider's schema is not aligned with the schema of the record batch produced by it, in particular column __delta_rs_path has type Utf8 according to the table provider, but in the record batches after scan it has type Dict(int32,Utf8), this error can be reproduced by

#[tokio::test]
async fn test_unnest_provider_with_mistmatch_schema() -> Result<()> {
    use datafusion::datasource::MemTable;
    // define a schema.
    let schema = Arc::new(Schema::new(vec![
        Field::new("column1", DataType::Utf8, false),
        Field::new(
            "arr_col",
            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
            true,
        ),
    ]));
    let strings = vec!["a", "b", "c"];
    let data = vec![
        Some(vec![Some(1), Some(2)]),
        Some(vec![Some(1), Some(2)]),
        Some(vec![Some(1), Some(2)]),
    ];
    let string_data = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
    let keys = Int32Array::from_iter(0..strings.len() as i32);

    let utf8_values = StringArray::from(strings.clone());
    let utf8_dict = DictionaryArray::new(keys.clone(), Arc::new(utf8_values.clone()));

    let batch1 = RecordBatch::try_new(
        schema.clone(),
        vec![Arc::new(utf8_dict.clone()), Arc::new(string_data)],
    )?;

    // declare a new context. In spark API, this corresponds to a new spark SQLsession
    let ctx = SessionContext::new();

    // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
    let provider = Arc::new(MemTable::try_new(schema, vec![vec![batch1]])?);
    ctx.register_table("t", provider.clone())?;
    let plan =
        LogicalPlanBuilder::scan("t", Arc::new(DefaultTableSource::new(provider)), None)?
            .unnest_column("arr_col")?
            .build()?;
    let df = ctx.execute_logical_plan(plan.clone()).await?;
    let ret = df.collect().await?;
    println!("{:?}", ret);
    Ok(())
}

I'm gathering more context here to check where we do this implicit type coercion/(or do we allow this type coercion at execution time instead of at plan time)

[duongcongtoai]

and i don't think this error is caused by unnest, for example with this new plan

    let plan =
        LogicalPlanBuilder::scan("t", Arc::new(DefaultTableSource::new(provider)), None)?
            .filter(binary_expr(
                col("column1"),
                datafusion_expr::Operator::Gt,
                lit(ScalarValue::new_utf8("some_val")),
            ))?
            // .unnest_column("arr_col")?
            .build()?;

It will throw the same error because datafusion cannot handle the type difference at execution time, meaning the table provider should correctly provide the type for column column1 for type_coercion to happen (at plan time)

[jayshrivastava]

I'm also getting this error when aggregating on a dictionary encoded column:

pub async fn register_temp_parquet_table(
    table_name: &str,
    schema: SchemaRef,
    batches: Vec<RecordBatch>,
    ctx: &SessionContext,
) -> Result<PathBuf> {
    if batches.is_empty() {
        return Err(datafusion::error::DataFusionError::Execution(
            "cannot create parquet file from empty batch list".to_string(),
        ));
    }
    for batch in &batches {
        if batch.schema() != schema {
            return Err(datafusion::error::DataFusionError::Execution(
                "all batches must have the same schema".to_string(),
            ));
        }
    }

    let temp_dir = std::env::temp_dir();
    let file_id = Uuid::new_v4();
    let temp_file_path = temp_dir.join(format!("{table_name}_{file_id}.parquet",));

    let file = std::fs::File::create(&temp_file_path)?;
    let schema = batches[0].schema();
    let mut writer = ArrowWriter::try_new(file, schema, None)?;

    for batch in batches {
        writer.write(&batch)?;
    }
    writer.close()?;

    ctx.register_parquet(
        table_name,
        temp_file_path.to_string_lossy().as_ref(),
        ParquetReadOptions::default(),
    )
    .await?;

    Ok(temp_file_path)
}

let config = SessionConfig::new().with_target_partitions(2);

let state = SessionStateBuilder::new()
    .with_default_features()
    .with_config(config)
    .build();

let ctx = SessionContext::from(state);

let schema2 = Arc::new(Schema::new(vec![
    Field::new("id", DataType::Int32, false),
    Field::new("name", DataType::Utf8, false),
    Field::new("phone", DataType::Utf8, false),
    Field::new("balance", DataType::Float64, false),
    Field::new(
        "company",
        DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
        false,
    ),
]));

let batches2 = vec![
    RecordBatch::try_new(
        schema2.clone(),
        vec![
            Arc::new(Int32Array::from(vec![1, 2, 3])),
            Arc::new(StringArray::from(vec![
                "customer1",
                "customer2",
                "customer3",
            ])),
            Arc::new(StringArray::from(vec![
                "13-123-4567",
                "31-456-7890",
                "23-789-0123",
            ])),
            Arc::new(datafusion::arrow::array::Float64Array::from(vec![
                100.5, 250.0, 50.25,
            ])),
            Arc::new(
                vec!["company1", "company1", "company1"]
                    .into_iter()
                    .collect::<arrow::array::DictionaryArray<UInt16Type>>(),
            ),
        ],
    )
    .unwrap(),
];

// Register the test data as parquet tables
let _ = register_temp_parquet_table("table1", schema1, batches1, &ctx)
    .await
    .unwrap();

The query SELECT distinct company from table2 fails with Shared(ArrowError(InvalidArgumentError("column types must match schema types, expected Dictionary(UInt16, Utf8) but found Utf8 at column index 0"), Some("")))

[jayshrivastava]

I rolled back my code to df 51 and the error does not seem to occur.