ensure MemTable has at least one partition (#16754)

Co-authored-by: Yongting You <[email protected]>

Author: waynexia

datafusion-examples/examples/sql_analysis.rs

@@ -274,7 +274,10 @@ from
     for table in tables {
         ctx.register_table(
             table.name,
-            Arc::new(MemTable::try_new(Arc::new(table.schema.clone()), vec![])?),
+            Arc::new(MemTable::try_new(
+                Arc::new(table.schema.clone()),
+                vec![vec![]],
+            )?),
         )?;
     }
     // We can create a LogicalPlan from a SQL query like this

datafusion/catalog/src/memory/table.rs

@@ -67,8 +67,16 @@ pub struct MemTable {
 }
 
 impl MemTable {
-    /// Create a new in-memory table from the provided schema and record batches
+    /// Create a new in-memory table from the provided schema and record batches.
+    ///
+    /// Requires at least one partition. To construct an empty `MemTable`, pass
+    /// `vec![vec![]]` as the `partitions` argument, this represents one partition with
+    /// no batches.
     pub fn try_new(schema: SchemaRef, partitions: Vec<Vec<RecordBatch>>) -> Result<Self> {
+        if partitions.is_empty() {
+            return plan_err!("No partitions provided, expected at least one partition");
+        }
+
         for batches in partitions.iter().flatten() {
             let batches_schema = batches.schema();
             if !schema.contains(&batches_schema) {

datafusion/core/src/datasource/memory_test.rs

@@ -446,7 +446,7 @@ mod tests {
             .unwrap_err();
         // Ensure that there is a descriptive error message
         assert_eq!(
-            "Error during planning: Cannot insert into MemTable with zero partitions",
+            "Error during planning: No partitions provided, expected at least one partition",
             experiment_result.strip_backtrace()
         );
         Ok(())

datafusion/core/tests/dataframe/mod.rs

@@ -4851,7 +4851,7 @@ async fn use_var_provider() -> Result<()> {
         Field::new("bar", DataType::Int64, false),
     ]));
 
-    let mem_table = Arc::new(MemTable::try_new(schema, vec![])?);
+    let mem_table = Arc::new(MemTable::try_new(schema, vec![vec![]])?);
 
     let config = SessionConfig::new()
         .with_target_partitions(4)

docs/source/library-user-guide/building-logical-plans.md

@@ -181,7 +181,7 @@ async fn main() -> Result<(), DataFusionError> {
     // TableProvider. For this example, we don't provide any data
     // but in production code, this would have `RecordBatch`es with
     // in memory data
-    let table_provider = Arc::new(MemTable::try_new(Arc::new(schema), vec![])?);
+    let table_provider = Arc::new(MemTable::try_new(Arc::new(schema), vec![vec![]])?);
     // Use the provider_as_source function to convert the TableProvider to a table source
     let table_source = provider_as_source(table_provider);