Add test_arrow_mysql_roundtrip

Author: gengteng

tests/mysql/common.rs

@@ -12,7 +12,7 @@ use crate::{
 const MYSQL_ROOT_PASSWORD: &str = "integration-test-pw";
 const MYSQL_DOCKER_CONTAINER: &str = "runtime-integration-test-mysql";
 
-fn get_mysql_params(port: usize) -> HashMap<String, SecretString> {
+pub(super) fn get_mysql_params(port: usize) -> HashMap<String, SecretString> {
     let mut params = HashMap::new();
     params.insert(
         "mysql_host".to_string(),

tests/mysql/mod.rs

@@ -1,3 +1,4 @@
+use crate::arrow_record_batch_gen::*;
 use datafusion::execution::context::SessionContext;
 use datafusion_table_providers::{
     mysql::DynMySQLConnectionPool, sql::sql_provider_datafusion::SqlTable,
@@ -9,8 +10,16 @@ use arrow::{
     array::*,
     datatypes::{i256, DataType, Field, Schema, TimeUnit, UInt16Type},
 };
-
+use arrow_schema::SchemaRef;
+use datafusion::catalog::TableProviderFactory;
+use datafusion::common::{Constraints, ToDFSchema};
+use datafusion::logical_expr::CreateExternalTable;
+use datafusion::physical_plan::collect;
+use datafusion::physical_plan::memory::MemoryExec;
+use datafusion_federation::schema_cast::record_convert::try_cast_to;
+use datafusion_table_providers::mysql::MySQLTableProviderFactory;
 use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection;
+use secrecy::ExposeSecret;
 
 use crate::docker::RunningContainer;
 
@@ -578,6 +587,73 @@ async fn arrow_mysql_one_way(
     record_batch
 }
 
+async fn arrow_mysql_round_trip(
+    port: usize,
+    arrow_record: RecordBatch,
+    source_schema: SchemaRef,
+    table_name: &str,
+) {
+    let factory = MySQLTableProviderFactory::new();
+    let ctx = SessionContext::new();
+    let cmd = CreateExternalTable {
+        schema: Arc::new(arrow_record.schema().to_dfschema().expect("to df schema")),
+        name: table_name.into(),
+        location: "".to_string(),
+        file_type: "".to_string(),
+        table_partition_cols: vec![],
+        if_not_exists: false,
+        definition: None,
+        order_exprs: vec![],
+        unbounded: false,
+        options: common::get_mysql_params(port)
+            .into_iter()
+            .map(|(k, v)| (k, v.expose_secret().to_string()))
+            .collect(),
+        constraints: Constraints::empty(),
+        column_defaults: Default::default(),
+    };
+    let table_provider = factory
+        .create(&ctx.state(), &cmd)
+        .await
+        .expect("table provider created");
+
+    let ctx = SessionContext::new();
+    let mem_exec = MemoryExec::try_new(&[vec![arrow_record.clone()]], arrow_record.schema(), None)
+        .expect("memory exec created");
+    let insert_plan = table_provider
+        .insert_into(&ctx.state(), Arc::new(mem_exec), true)
+        .await
+        .expect("insert plan created");
+
+    let _ = collect(insert_plan, ctx.task_ctx())
+        .await
+        .expect("insert done");
+    ctx.register_table(table_name, table_provider)
+        .expect("Table should be registered");
+    let sql = format!("SELECT * FROM {table_name}");
+    let df = ctx
+        .sql(&sql)
+        .await
+        .expect("DataFrame should be created from query");
+
+    let record_batch = df.collect().await.expect("RecordBatch should be collected");
+
+    tracing::debug!("Original Arrow Record Batch: {:?}", arrow_record.columns());
+    tracing::debug!(
+        "MySQL returned Record Batch: {:?}",
+        record_batch[0].columns()
+    );
+
+    let casted_result =
+        try_cast_to(record_batch[0].clone(), source_schema).expect("Failed to cast record batch");
+
+    // Check results
+    assert_eq!(record_batch.len(), 1);
+    assert_eq!(record_batch[0].num_rows(), arrow_record.num_rows());
+    assert_eq!(record_batch[0].num_columns(), arrow_record.num_columns());
+    assert_eq!(arrow_record, casted_result);
+}
+
 async fn start_mysql_container(port: usize) -> RunningContainer {
     let running_container = common::start_mysql_docker_container(port)
         .await
@@ -588,6 +664,34 @@ async fn start_mysql_container(port: usize) -> RunningContainer {
     running_container
 }
 
+#[rstest]
+#[case::binary(get_arrow_binary_record_batch(), "binary")]
+#[case::int(get_arrow_int_record_batch(), "int")]
+#[case::float(get_arrow_float_record_batch(), "float")]
+#[case::utf8(get_arrow_utf8_record_batch(), "utf8")]
+#[case::time(get_arrow_time_record_batch(), "time")]
+#[case::timestamp(get_arrow_timestamp_record_batch(), "timestamp")]
+#[case::date(get_arrow_date_record_batch(), "date")]
+#[case::struct_type(get_arrow_struct_record_batch(), "struct")]
+#[case::decimal(get_arrow_decimal_record_batch(), "decimal")]
+#[case::interval(get_arrow_interval_record_batch(), "interval")]
+#[case::duration(get_arrow_duration_record_batch(), "duration")]
+#[case::list(get_arrow_list_record_batch(), "list")]
+#[case::null(get_arrow_null_record_batch(), "null")]
+#[case::bytea_array(get_arrow_bytea_array_record_batch(), "bytea_array")]
+#[test_log::test(tokio::test)]
+async fn test_arrow_mysql_roundtrip(
+    #[case] arrow_result: (RecordBatch, SchemaRef),
+    #[case] table_name: &str,
+) {
+    let port = crate::get_random_port();
+    let mysql_container = start_mysql_container(port).await;
+
+    arrow_mysql_round_trip(port, arrow_result.0, arrow_result.1, table_name).await;
+
+    mysql_container.remove().await.expect("container to stop");
+}
+
 #[rstest]
 #[test_log::test(tokio::test)]
 async fn test_mysql_arrow_oneway() {