Update tests

Author: alamb

datafusion/core/src/datasource/memory.rs

@@ -228,8 +228,8 @@ mod tests {
     use crate::from_slice::FromSlice;
     use crate::physical_plan::collect;
     use crate::prelude::SessionContext;
-    use arrow::array::Int32Array;
-    use arrow::datatypes::{DataType, Field, Schema};
+    use arrow::array::{AsArray, Int32Array};
+    use arrow::datatypes::{DataType, Field, Schema, UInt64Type};
     use arrow::error::ArrowError;
     use datafusion_expr::LogicalPlanBuilder;
     use futures::StreamExt;
@@ -467,6 +467,11 @@ mod tests {
         initial_data: Vec<Vec<RecordBatch>>,
         inserted_data: Vec<Vec<RecordBatch>>,
     ) -> Result<Vec<Vec<RecordBatch>>> {
+        let expected_count: u64 = inserted_data
+            .iter()
+            .flat_map(|batches| batches.iter().map(|batch| batch.num_rows() as u64))
+            .sum();
+
         // Create a new session context
         let session_ctx = SessionContext::new();
         // Create and register the initial table with the provided schema and data
@@ -490,8 +495,8 @@ mod tests {
 
         // Execute the physical plan and collect the results
         let res = collect(plan, session_ctx.task_ctx()).await?;
-        // Ensure the result is empty after the insert operation
-        assert!(res.is_empty());
+        assert_eq!(extract_count(res), expected_count);
+
         // Read the data from the initial table and store it in a vector of partitions
         let mut partitions = vec![];
         for partition in initial_table.batches.iter() {
@@ -501,6 +506,32 @@ mod tests {
         Ok(partitions)
     }
 
+    /// Returns the value of results:
+    ///
+    /// "+-------+",
+    /// "| count |",
+    /// "+-------+",
+    /// "| 6     |",
+    /// "+-------+",
+    fn extract_count(res: Vec<RecordBatch>) -> u64 {
+        assert_eq!(res.len(), 1, "expected one batch, got {}", res.len());
+        let batch = &res[0];
+        assert_eq!(
+            batch.num_columns(),
+            1,
+            "expected 1 column, got {}",
+            batch.num_columns()
+        );
+        let col = batch.column(0).as_primitive::<UInt64Type>();
+        assert_eq!(col.len(), 1, "expected 1 row, got {}", col.len());
+        let val = col
+            .iter()
+            .next()
+            .expect("had value")
+            .expect("expected non null");
+        val
+    }
+
     // Test inserting a single batch of data into a single partition
     #[tokio::test]
     async fn test_insert_into_single_partition() -> Result<()> {

datafusion/core/src/physical_plan/memory.rs

@@ -238,133 +238,3 @@ impl RecordBatchStream for MemoryStream {
         self.schema.clone()
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::datasource::streaming::PartitionStream;
-    use crate::datasource::{MemTable, TableProvider};
-    use crate::from_slice::FromSlice;
-    use crate::physical_plan::stream::RecordBatchStreamAdapter;
-    use crate::physical_plan::streaming::StreamingTableExec;
-    use crate::physical_plan::ColumnStatistics;
-    use crate::physical_plan::{collect, displayable, SendableRecordBatchStream};
-    use crate::prelude::{CsvReadOptions, SessionContext};
-    use crate::test_util;
-    use arrow::array::Int32Array;
-    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-    use arrow::record_batch::RecordBatch;
-    use datafusion_common::Result;
-    use datafusion_execution::config::SessionConfig;
-    use datafusion_execution::TaskContext;
-    use futures::StreamExt;
-    use std::sync::Arc;
-
-    fn mock_data() -> Result<(SchemaRef, RecordBatch)> {
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("a", DataType::Int32, false),
-            Field::new("b", DataType::Int32, false),
-            Field::new("c", DataType::Int32, true),
-            Field::new("d", DataType::Int32, false),
-        ]));
-
-        let batch = RecordBatch::try_new(
-            schema.clone(),
-            vec![
-                Arc::new(Int32Array::from_slice([1, 2, 3])),
-                Arc::new(Int32Array::from_slice([4, 5, 6])),
-                Arc::new(Int32Array::from(vec![None, None, Some(9)])),
-                Arc::new(Int32Array::from_slice([7, 8, 9])),
-            ],
-        )?;
-
-        Ok((schema, batch))
-    }
-
-    #[tokio::test]
-    async fn test_with_projection() -> Result<()> {
-        let session_ctx = SessionContext::new();
-        let task_ctx = session_ctx.task_ctx();
-        let (schema, batch) = mock_data()?;
-
-        let executor = MemoryExec::try_new(&[vec![batch]], schema, Some(vec![2, 1]))?;
-        let statistics = executor.statistics();
-
-        assert_eq!(statistics.num_rows, Some(3));
-        assert_eq!(
-            statistics.column_statistics,
-            Some(vec![
-                ColumnStatistics {
-                    null_count: Some(2),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-                ColumnStatistics {
-                    null_count: Some(0),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-            ])
-        );
-
-        // scan with projection
-        let mut it = executor.execute(0, task_ctx)?;
-        let batch2 = it.next().await.unwrap()?;
-        assert_eq!(2, batch2.schema().fields().len());
-        assert_eq!("c", batch2.schema().field(0).name());
-        assert_eq!("b", batch2.schema().field(1).name());
-        assert_eq!(2, batch2.num_columns());
-
-        Ok(())
-    }
-
-    #[tokio::test]
-    async fn test_without_projection() -> Result<()> {
-        let session_ctx = SessionContext::new();
-        let task_ctx = session_ctx.task_ctx();
-        let (schema, batch) = mock_data()?;
-
-        let executor = MemoryExec::try_new(&[vec![batch]], schema, None)?;
-        let statistics = executor.statistics();
-
-        assert_eq!(statistics.num_rows, Some(3));
-        assert_eq!(
-            statistics.column_statistics,
-            Some(vec![
-                ColumnStatistics {
-                    null_count: Some(0),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-                ColumnStatistics {
-                    null_count: Some(0),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-                ColumnStatistics {
-                    null_count: Some(2),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-                ColumnStatistics {
-                    null_count: Some(0),
-                    max_value: None,
-                    min_value: None,
-                    distinct_count: None,
-                },
-            ])
-        );
-
-        let mut it = executor.execute(0, task_ctx)?;
-        let batch1 = it.next().await.unwrap()?;
-        assert_eq!(4, batch1.schema().fields().len());
-        assert_eq!(4, batch1.num_columns());
-
-        Ok(())
-    }
-}