From f1d95cd0b11644f2a424d8830a08b3652fa41f11 Mon Sep 17 00:00:00 2001
From: Matthew Gapp <61894094+matthewgapp@users.noreply.github.com>
Date: Sun, 24 Sep 2023 12:13:48 -0700
Subject: [PATCH] take schema from record batch when creating mem table if any

---
 datafusion/core/src/execution/context.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/datafusion/core/src/execution/context.rs b/datafusion/core/src/execution/context.rs
index 78b79680b4d7..31b6f55b2946 100644
--- a/datafusion/core/src/execution/context.rs
+++ b/datafusion/core/src/execution/context.rs
@@ -534,6 +534,13 @@ impl SessionContext {
                 let physical = DataFrame::new(self.state(), input);
 
                 let batches: Vec<_> = physical.collect_partitioned().await?;
+                let schema = {
+                    if let Some(batch) = batches.first().and_then(|b| b.first()) {
+                        batch.schema().clone()
+                    } else {
+                        schema
+                    }
+                };
                 let table = Arc::new(
                     // pass constraints to the mem table.
                     MemTable::try_new(schema, batches)?.with_constraints(constraints),