Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: casting when data to be written does not match table schema #1427

Merged
merged 4 commits into from
Jun 3, 2023
Merged

fix: casting when data to be written does not match table schema #1427

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
36ba8be
Add test that fails when different datatypes are written
Blajda Jun 3, 2023
9b95848
Use table schema if it exists when casting batches
Blajda Jun 3, 2023
d1d0a25
Fix casting when types do not match table schema
Blajda Jun 3, 2023
ba83dbd
Allow users to specify cast options on write
Blajda Jun 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 7 additions & 17 deletions rust/src/operations/delete.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,6 +681,7 @@ mod tests {

use crate::action::*;
use crate::operations::DeltaOps;
use crate::writer::test_utils::datafusion::get_data;
use crate::writer::test_utils::{get_arrow_schema, get_delta_schema};
use crate::DeltaTable;
use arrow::array::Int32Array;
Expand All @@ -704,17 +705,6 @@ mod tests {
table
}

async fn get_data(table: DeltaTable) -> Vec<RecordBatch> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is a good refactor 👍

let ctx = SessionContext::new();
ctx.register_table("test", Arc::new(table)).unwrap();
ctx.sql("select * from test")
.await
.unwrap()
.collect()
.await
.unwrap()
}

#[tokio::test]
async fn test_delete_default() {
let schema = get_arrow_schema(&None);
Expand Down Expand Up @@ -850,7 +840,7 @@ mod tests {
"+----+-------+------------+",
];

let actual = get_data(table).await;
let actual = get_data(&table).await;
assert_batches_sorted_eq!(&expected, &actual);
}

Expand Down Expand Up @@ -898,7 +888,7 @@ mod tests {
"| 2 |",
"+-------+",
];
let actual = get_data(table).await;
let actual = get_data(&table).await;
assert_batches_sorted_eq!(&expected, &actual);

// Validate behaviour of less than
Expand All @@ -919,7 +909,7 @@ mod tests {
"| 4 |",
"+-------+",
];
let actual = get_data(table).await;
let actual = get_data(&table).await;
assert_batches_sorted_eq!(&expected, &actual);

// Validate behaviour of less plus not null
Expand All @@ -938,7 +928,7 @@ mod tests {
"| 4 |",
"+-------+",
];
let actual = get_data(table).await;
let actual = get_data(&table).await;
assert_batches_sorted_eq!(&expected, &actual);
}

Expand Down Expand Up @@ -997,7 +987,7 @@ mod tests {
"+----+-------+------------+",
];

let actual = get_data(table).await;
let actual = get_data(&table).await;
assert_batches_sorted_eq!(&expected, &actual);
}

Expand Down Expand Up @@ -1058,7 +1048,7 @@ mod tests {
"| B | 20 | 2021-02-03 |",
"+----+-------+------------+",
];
let actual = get_data(table).await;
let actual = get_data(&table).await;
assert_batches_sorted_eq!(&expected, &actual);
}

Expand Down
63 changes: 60 additions & 3 deletions rust/src/operations/write.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,30 +221,34 @@ pub(crate) async fn write_execution_plan(
.current_metadata()
.and_then(|meta| meta.schema.get_invariants().ok())
.unwrap_or_default();

// Use input schema to prevent wrapping partitions columns into a dictionary.
let schema = snapshot.input_schema().unwrap_or(plan.schema());

let checker = DeltaDataChecker::new(invariants);

// Write data to disk
let mut tasks = vec![];
for i in 0..plan.output_partitioning().partition_count() {
let inner_plan = plan.clone();
let inner_schema = schema.clone();
let task_ctx = Arc::new(TaskContext::from(&state));
let config = WriterConfig::new(
inner_plan.schema(),
inner_schema.clone(),
partition_columns.clone(),
writer_properties.clone(),
target_file_size,
write_batch_size,
);
let mut writer = DeltaWriter::new(object_store.clone(), config);
let checker_stream = checker.clone();
let schema = inner_plan.schema().clone();
let mut stream = inner_plan.execute(i, task_ctx)?;
let handle: tokio::task::JoinHandle<DeltaResult<Vec<Add>>> =
tokio::task::spawn(async move {
while let Some(maybe_batch) = stream.next().await {
let batch = maybe_batch?;
checker_stream.check_batch(&batch).await?;
let arr = cast_record_batch(&batch, schema.clone())?;
let arr = cast_record_batch(&batch, inner_schema.clone())?;
writer.write(&arr).await?;
}
writer.close().await
Expand Down Expand Up @@ -483,7 +487,12 @@ fn cast_record_batch(
mod tests {
use super::*;
use crate::operations::DeltaOps;
use crate::writer::test_utils::datafusion::get_data;
use crate::writer::test_utils::{get_delta_schema, get_record_batch};
use arrow::datatypes::Field;
use arrow::datatypes::Schema as ArrowSchema;
use arrow_array::{Int32Array, StringArray};
use datafusion::assert_batches_sorted_eq;
use serde_json::json;

#[tokio::test]
Expand Down Expand Up @@ -526,6 +535,54 @@ mod tests {
assert_eq!(table.get_file_uris().count(), 1)
}

#[tokio::test]
async fn test_write_different_types() {
// Ensure write data is casted when data of a different type from the table is provided.
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"value",
arrow::datatypes::DataType::Int32,
true,
)]));

let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(Int32Array::from(vec![Some(0), None]))],
)
.unwrap();
let table = DeltaOps::new_in_memory().write(vec![batch]).await.unwrap();

let schema = Arc::new(ArrowSchema::new(vec![Field::new(
"value",
arrow::datatypes::DataType::Utf8,
true,
)]));

let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(StringArray::from(vec![
Some("Test123".to_owned()),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm confused, where does this value go in the expected table? 😕

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right below, in line 570?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since it cannot be parsed as an int it will result in a null value. Which aligns with ansi sql but maybe we should only allow that if the user opts in?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh nevermind. It seems like this is being inserted as null? That doesn't seem like what a user would want.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I think we could disable that by passing safe: true in CastOptions: https://docs.rs/arrow-cast/40.0.0/arrow_cast/cast/struct.CastOptions.html

Some("123".to_owned()),
None,
]))],
)
.unwrap();

let table = DeltaOps::from(table).write(vec![batch]).await.unwrap();
let expected = [
"+-------+",
"| value |",
"+-------+",
"| |",
"| |",
"| |",
"| 123 |",
"| 0 |",
"+-------+",
];
let actual = get_data(&table).await;
assert_batches_sorted_eq!(&expected, &actual);
}

#[tokio::test]
async fn test_write_nonexistent() {
let batch = get_record_batch(None, false);
Expand Down
20 changes: 20 additions & 0 deletions rust/src/writer/test_utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,3 +211,23 @@ pub async fn create_initialized_table(partition_cols: &[String]) -> DeltaTable {

table
}

#[cfg(feature = "datafusion")]
pub mod datafusion {
use crate::DeltaTable;
use arrow_array::RecordBatch;
use datafusion::prelude::SessionContext;
use std::sync::Arc;

pub async fn get_data(table: &DeltaTable) -> Vec<RecordBatch> {
let table = DeltaTable::new_with_state(table.object_store(), table.state.clone());
let ctx = SessionContext::new();
ctx.register_table("test", Arc::new(table)).unwrap();
ctx.sql("select * from test")
.await
.unwrap()
.collect()
.await
.unwrap()
}
}