| 37 | |
| 38 | #[tokio::test] |
| 39 | async fn test_with_projection() -> Result<()> { |
| 40 | let session_ctx = SessionContext::new(); |
| 41 | let task_ctx = session_ctx.task_ctx(); |
| 42 | let schema = Arc::new(Schema::new(vec![ |
| 43 | Field::new("a", DataType::Int32, false), |
| 44 | Field::new("b", DataType::Int32, false), |
| 45 | Field::new("c", DataType::Int32, false), |
| 46 | Field::new("d", DataType::Int32, true), |
| 47 | ])); |
| 48 | |
| 49 | let batch = RecordBatch::try_new( |
| 50 | schema.clone(), |
| 51 | vec![ |
| 52 | Arc::new(Int32Array::from(vec![1, 2, 3])), |
| 53 | Arc::new(Int32Array::from(vec![4, 5, 6])), |
| 54 | Arc::new(Int32Array::from(vec![7, 8, 9])), |
| 55 | Arc::new(Int32Array::from(vec![None, None, Some(9)])), |
| 56 | ], |
| 57 | )?; |
| 58 | |
| 59 | let provider = MemTable::try_new(schema, vec![vec![batch]])?; |
| 60 | |
| 61 | // scan with projection |
| 62 | let exec = provider |
| 63 | .scan(&session_ctx.state(), Some(&vec![2, 1]), &[], None) |
| 64 | .await?; |
| 65 | |
| 66 | let mut it = exec.execute(0, task_ctx)?; |
| 67 | let batch2 = it.next().await.unwrap()?; |
| 68 | assert_eq!(2, batch2.schema().fields().len()); |
| 69 | assert_eq!("c", batch2.schema().field(0).name()); |
| 70 | assert_eq!("b", batch2.schema().field(1).name()); |
| 71 | assert_eq!(2, batch2.num_columns()); |
| 72 | |
| 73 | Ok(()) |
| 74 | } |
| 75 | |
| 76 | #[tokio::test] |
| 77 | async fn test_without_projection() -> Result<()> { |