Skip to content

Commit 6e84421

Browse files
committed
ARROW-8869: [Rust] [DataFusion] Add support for new scan nodes to type coercion rule
Add support for new scan nodes to type coercion rule. I need to add unit tests. Closes #7230 from andygrove/ARROW-8869 Authored-by: Andy Grove <andygrove73@gmail.com> Signed-off-by: Andy Grove <andygrove73@gmail.com>
1 parent 3533e01 commit 6e84421

File tree

1 file changed

+27
-1
lines changed

1 file changed

+27
-1
lines changed

rust/datafusion/src/optimizer/type_coercion.rs

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,9 @@ impl<'a> OptimizerRule for TypeCoercionRule<'a> {
167167
)?
168168
.build(),
169169
LogicalPlan::TableScan { .. } => Ok(plan.clone()),
170+
LogicalPlan::InMemoryScan { .. } => Ok(plan.clone()),
171+
LogicalPlan::ParquetScan { .. } => Ok(plan.clone()),
172+
LogicalPlan::CsvScan { .. } => Ok(plan.clone()),
170173
LogicalPlan::EmptyRelation { .. } => Ok(plan.clone()),
171174
LogicalPlan::Limit { .. } => Ok(plan.clone()),
172175
LogicalPlan::CreateExternalTable { .. } => Ok(plan.clone()),
@@ -182,10 +185,33 @@ impl<'a> OptimizerRule for TypeCoercionRule<'a> {
182185
mod tests {
183186
use super::*;
184187
use crate::execution::context::ExecutionContext;
188+
use crate::execution::physical_plan::csv::CsvReadOptions;
185189
use crate::logicalplan::Expr::*;
186-
use crate::logicalplan::Operator;
190+
use crate::logicalplan::{col, Operator};
191+
use crate::test::arrow_testdata_path;
187192
use arrow::datatypes::{DataType, Field, Schema};
188193

194+
#[test]
195+
fn test_with_csv_plan() -> Result<()> {
196+
let testdata = arrow_testdata_path();
197+
let path = format!("{}/csv/aggregate_test_100.csv", testdata);
198+
199+
let options = CsvReadOptions::new().schema_infer_max_records(100);
200+
let plan = LogicalPlanBuilder::scan_csv(&path, options, None)?
201+
.filter(col("c7").lt(&col("c12")))?
202+
.build()?;
203+
204+
let scalar_functions = HashMap::new();
205+
let mut rule = TypeCoercionRule::new(&scalar_functions);
206+
let plan = rule.optimize(&plan)?;
207+
208+
assert!(
209+
format!("{:?}", plan).starts_with("Selection: CAST(#c7 AS Float64) Lt #c12")
210+
);
211+
212+
Ok(())
213+
}
214+
189215
#[test]
190216
fn test_add_i32_i64() {
191217
binary_cast_test(

0 commit comments

Comments
 (0)