DataSink::write_all given invalid RecordBatchStream

[gatesn]

Describe the bug

We were trying to implement a DataSink and found that we were being given different schemas for the record batches than per the RecordBatchStream.

STREAM DTYPE Schema { fields: [Field { name: "c1", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "c2", data_type: Int32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }
RB SCHEMA: Schema { fields: [Field { name: "c1", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "c2", data_type: Int32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }

STREAM DTYPE Schema { fields: [Field { name: "c1", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "c2", data_type: Int32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }
RB SCHEMA: Schema { fields: [Field { name: "c1", data_type: Utf8, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "c2", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }

The first (correct) invocation comes from assembling a logical plan with LogicalPlanBuilder::insert_into

The second (incorrect) invocation comes from session.sql("INSERT INTO my_tbl VALUES ('hello', 42::INT);")

I figured I'd add an assertion into the RecordBatchStreamAdapter, and it looks like ~12 tests fail on main right now with mismatched schemas. I wonder if it's worth adding that as a debug assertion?

https://github.com/gatesn/datafusion/pull/new/ngates/record-batch-stream-schema

cc @AdamGS

To Reproduce

No response

Expected behavior

No response

Additional context

No response

[ozankabak]

Is this a bug in DataSink code, or is it a downstream bug that becomes somewhat hard to notice because there is no debug_assert?

[gatesn]

I imagine it's downstream, and of course, the debug assert only catches bad RecordBatchStream impls that used the Adapter.

This specific bug may not even be caught by the debug assert if it doesn't use the Adapter somewhere.

[zhuqi-lucas]

Thank you @gatesn for the report.

Can you provide the full code or sql to reproduce the it? So we can solve it more quickly.

[gatesn]

Yes, this code fails the schema assertion:

https://github.com/apache/datafusion/compare/main...gatesn:datafusion:ngates/record-batch-stream-schema?expand=1#diff-9b1672adeba35025e24d21f8d7da2f0e87487231775c5e470edf56886f90c837R811-R833

[alamb]

I wonder if it's worth adding that as a debug assertion?

I agree this would be good to add -- most similar check in DataFusion do a check and raise an DataFusionError;:Internal if hit

[ozankabak]

The only thing I would be worried about is the potential overhead. I get the feeling that this is something that should be done "once" but I'm not sure how we can do it. Maybe utilize the invariant and sanity checkers to check the contract somehow. Let's think about this.

[zhuqi-lucas]

The second (incorrect) invocation comes from session.sql("INSERT INTO my_tbl VALUES ('hello', 42::INT);")

          LogicalPlan::Dml(DmlStatement {
                table_name,
                op: WriteOp::Insert(insert_op),
                ..
            }) => {
                let name = table_name.table();
                let schema = session_state.schema_for_ref(table_name.clone())?;
                if let Some(provider) = schema.table(name).await? {
                    let input_exec = children.one()?;
                    provider
                        .insert_into(session_state, input_exec, *insert_op)
                        .await?
                } else {
                    return exec_err!("Table '{table_name}' does not exist");
                }
            }

I think we may can add check during the physical plan generation, we have the table schema, also we have the input_exec to insert into, so we can check here?

[gatesn]

Is this overhead a problem if it's only a debug_assertion? It will be compiled out in release builds. Or are you looking for a release assertion to put somewhere?

[rkrishn7]

Poked around and I think the specific issue here is due to the fact that the schema assigned to LogicalPlan::Values during planning defaults all its fields to nullable.

I've opened a PR to capture nullability from the table schema for the specified columns.

[jonahgao]

I am concerned whether union will also lead to similar behavior when the nullability of the two inputs is different.