test_utils: add utility to create test parquet tables from RecordBatch

This change adds a module `session_context` to `test_utils` which allows you to
- create temp parquet files from RecordBatch
- register this data in the `SessionContext` under a table name

This is handy for unit tests where you typically want to execute a distributed SQL
query on some small hard-coded `RecordBatch` data

Author: jayshrivastava

src/test_utils/mod.rs

@@ -3,4 +3,5 @@ pub mod insta;
 pub mod localhost;
 pub mod mock_exec;
 pub mod parquet;
+pub mod session_context;
 pub mod tpch;

src/test_utils/session_context.rs

@@ -0,0 +1,99 @@
+use arrow::record_batch::RecordBatch;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::error::Result;
+use datafusion::execution::context::SessionContext;
+use datafusion::prelude::ParquetReadOptions;
+use parquet::arrow::ArrowWriter;
+use std::path::PathBuf;
+use uuid::Uuid;
+
+/// Creates a temporary Parquet file from RecordBatches and registers it with the SessionContext
+/// under the provided table name. Returns the file name.
+///
+/// TODO: consider expanding this to support partitioned data
+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!("{}_{}.parquet", table_name, file_id,));
+
+    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)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{Int32Array, StringArray};
+    use arrow::datatypes::{DataType, Field, Schema};
+    use tokio::fs::remove_file;
+
+    use std::sync::Arc;
+
+    #[tokio::test]
+    async fn test_register_temp_parquet_table() {
+        let ctx = SessionContext::new();
+
+        // Create test data
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("name", DataType::Utf8, false),
+        ]));
+
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(Int32Array::from(vec![1, 2, 3])),
+                Arc::new(StringArray::from(vec!["a", "b", "c"])),
+            ],
+        )
+        .unwrap();
+
+        // Register temp table
+        let temp_file =
+            register_temp_parquet_table("test_table", schema.clone(), vec![batch], &ctx)
+                .await
+                .unwrap();
+
+        // Verify we can query it
+        let df = ctx.sql("SELECT COUNT(*) FROM test_table").await.unwrap();
+        let results = df.collect().await.unwrap();
+
+        assert_eq!(results.len(), 1);
+        assert_eq!(results[0].num_rows(), 1);
+
+        let _ = remove_file(temp_file).await;
+    }
+}