|
20 | 20 |
|
21 | 21 | use arrow::array::Int64Array;
|
22 | 22 | use arrow::datatypes::{DataType, Field};
|
| 23 | +use arrow_schema::Schema; |
| 24 | +use datafusion::datasource::{provider_as_source, ViewTable}; |
23 | 25 | use datafusion::execution::session_state::SessionStateBuilder;
|
24 |
| -use datafusion_common::{Column, DFSchema, Result, ScalarValue, Spans}; |
| 26 | +use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, ScalarValue, Spans}; |
25 | 27 | use datafusion_execution::TaskContext;
|
26 | 28 | use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
|
27 | 29 | use datafusion_expr::logical_plan::{LogicalPlan, Values};
|
28 |
| -use datafusion_expr::{Aggregate, AggregateUDF, Expr}; |
| 30 | +use datafusion_expr::{Aggregate, AggregateUDF, EmptyRelation, Expr, LogicalPlanBuilder, UNNAMED_TABLE}; |
29 | 31 | use datafusion_functions_aggregate::count::Count;
|
30 | 32 | use datafusion_physical_plan::collect;
|
31 | 33 | use std::collections::HashMap;
|
|
96 | 98 | };
|
97 | 99 | element
|
98 | 100 | }
|
| 101 | + |
| 102 | +#[test] |
| 103 | +fn inline_scan_projection_test() -> Result<()> { |
| 104 | + let name = UNNAMED_TABLE; |
| 105 | + let column = "a"; |
| 106 | + |
| 107 | + let schema = Schema::new(vec![ |
| 108 | + Field::new("a", DataType::Int32, false), |
| 109 | + Field::new("b", DataType::Int32, false), |
| 110 | + ]); |
| 111 | + let projection = vec![schema.index_of(column)?]; |
| 112 | + |
| 113 | + let provider = ViewTable::new( |
| 114 | + LogicalPlan::EmptyRelation(EmptyRelation { |
| 115 | + produce_one_row: false, |
| 116 | + schema: DFSchemaRef::new(DFSchema::try_from(schema)?), |
| 117 | + }), |
| 118 | + None, |
| 119 | + ); |
| 120 | + let source = provider_as_source(Arc::new(provider)); |
| 121 | + |
| 122 | + let plan = LogicalPlanBuilder::scan(name, source, Some(projection))?.build()?; |
| 123 | + assert_eq!( |
| 124 | + format!("{plan}"), |
| 125 | + format!("TableScan: {name} projection=[{column}]") |
| 126 | + ); |
| 127 | + |
| 128 | + Ok(()) |
| 129 | +} |
0 commit comments