Skip to content

Commit af2cda9

Browse files
haohuaijinDandandan
haohuaijin
and
Dandandan
authored
support scan empty projection (apache#7920)
* fix * modify push_down_filter for empty projection * clean code * fix ci * fix * Update datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs Co-authored-by: Daniël Heres <danielheres@gmail.com> * remove LogicalPlan::Values in push_down_projection --------- Co-authored-by: Daniël Heres <danielheres@gmail.com>
1 parent 8068d7f commit af2cda9

8 files changed

Lines changed: 31 additions & 181 deletions

File tree

‎datafusion/core/src/datasource/physical_plan/file_scan_config.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ use crate::{
3636
use arrow::array::{ArrayData, BufferBuilder};
3737
use arrow::buffer::Buffer;
3838
use arrow::datatypes::{ArrowNativeType, UInt16Type};
39-
use arrow_array::{ArrayRef, DictionaryArray, RecordBatch};
39+
use arrow_array::{ArrayRef, DictionaryArray, RecordBatch, RecordBatchOptions};
4040
use arrow_schema::{DataType, Field, Schema, SchemaRef};
4141
use datafusion_common::stats::Precision;
4242
use datafusion_common::{exec_err, ColumnStatistics, Statistics};
@@ -339,7 +339,13 @@ impl PartitionColumnProjector {
339339
),
340340
)
341341
}
342-
RecordBatch::try_new(Arc::clone(&self.projected_schema), cols).map_err(Into::into)
342+
343+
RecordBatch::try_new_with_options(
344+
Arc::clone(&self.projected_schema),
345+
cols,
346+
&RecordBatchOptions::new().with_row_count(Some(file_batch.num_rows())),
347+
)
348+
.map_err(Into::into)
343349
}
344350
}
345351

‎datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -149,10 +149,12 @@ impl TableProvider for CustomProvider {
149149
async fn scan(
150150
&self,
151151
_state: &SessionState,
152-
_: Option<&Vec<usize>>,
152+
projection: Option<&Vec<usize>>,
153153
filters: &[Expr],
154154
_: Option<usize>,
155155
) -> Result<Arc<dyn ExecutionPlan>> {
156+
let empty = Vec::new();
157+
let projection = projection.unwrap_or(&empty);
156158
match &filters[0] {
157159
Expr::BinaryExpr(BinaryExpr { right, .. }) => {
158160
let int_value = match &**right {
@@ -182,7 +184,10 @@ impl TableProvider for CustomProvider {
182184
};
183185

184186
Ok(Arc::new(CustomPlan {
185-
schema: self.zero_batch.schema(),
187+
schema: match projection.is_empty() {
188+
true => Arc::new(Schema::empty()),
189+
false => self.zero_batch.schema(),
190+
},
186191
batches: match int_value {
187192
0 => vec![self.zero_batch.clone()],
188193
1 => vec![self.one_batch.clone()],
@@ -191,7 +196,10 @@ impl TableProvider for CustomProvider {
191196
}))
192197
}
193198
_ => Ok(Arc::new(CustomPlan {
194-
schema: self.zero_batch.schema(),
199+
schema: match projection.is_empty() {
200+
true => Arc::new(Schema::empty()),
201+
false => self.zero_batch.schema(),
202+
},
195203
batches: vec![],
196204
})),
197205
}

‎datafusion/core/tests/sql/explain_analyze.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -788,7 +788,7 @@ async fn explain_logical_plan_only() {
788788
"logical_plan",
789789
"Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1)) AS COUNT(*)]]\
790790
\n SubqueryAlias: t\
791-
\n Projection: column2\
791+
\n Projection: \
792792
\n Values: (Utf8(\"a\"), Int64(1), Int64(100)), (Utf8(\"a\"), Int64(2), Int64(150))"
793793
]];
794794
assert_eq!(expected, actual);

‎datafusion/optimizer/src/push_down_projection.rs‎

Lines changed: 1 addition & 167 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ use crate::merge_projection::merge_projection;
2323
use crate::optimizer::ApplyOrder;
2424
use crate::push_down_filter::replace_cols_by_name;
2525
use crate::{OptimizerConfig, OptimizerRule};
26-
use arrow::datatypes::DataType;
2726
use arrow::error::Result as ArrowResult;
2827
use datafusion_common::ScalarValue::UInt8;
2928
use datafusion_common::{
@@ -149,10 +148,6 @@ impl OptimizerRule for PushDownProjection {
149148
{
150149
let mut used_columns: HashSet<Column> = HashSet::new();
151150
if projection_is_empty {
152-
let field = find_small_field(scan.projected_schema.fields()).ok_or(
153-
DataFusionError::Internal("Scan with empty schema".to_string()),
154-
)?;
155-
used_columns.insert(field.qualified_column());
156151
push_down_scan(&used_columns, scan, true)?
157152
} else {
158153
for expr in projection.expr.iter() {
@@ -163,17 +158,6 @@ impl OptimizerRule for PushDownProjection {
163158
plan.with_new_inputs(&[new_scan])?
164159
}
165160
}
166-
LogicalPlan::Values(values) if projection_is_empty => {
167-
let field = find_small_field(values.schema.fields()).ok_or(
168-
DataFusionError::Internal("Values with empty schema".to_string()),
169-
)?;
170-
let column = Expr::Column(field.qualified_column());
171-
172-
LogicalPlan::Projection(Projection::try_new(
173-
vec![column],
174-
Arc::new(child_plan.clone()),
175-
)?)
176-
}
177161
LogicalPlan::Union(union) => {
178162
let mut required_columns = HashSet::new();
179163
exprlist_to_columns(&projection.expr, &mut required_columns)?;
@@ -429,87 +413,6 @@ pub fn collect_projection_expr(projection: &Projection) -> HashMap<String, Expr>
429413
.collect::<HashMap<_, _>>()
430414
}
431415

432-
/// Accumulate the memory size of a data type measured in bits.
433-
///
434-
/// Types with a variable size get assigned with a fixed size which is greater than most
435-
/// primitive types.
436-
///
437-
/// While traversing nested types, `nesting` is incremented on every level.
438-
fn nested_size(data_type: &DataType, nesting: &mut usize) -> usize {
439-
use DataType::*;
440-
if data_type.is_primitive() {
441-
return data_type.primitive_width().unwrap_or(1) * 8;
442-
}
443-
444-
if data_type.is_nested() {
445-
*nesting += 1;
446-
}
447-
448-
match data_type {
449-
Null => 0,
450-
Boolean => 1,
451-
Binary | Utf8 => 128,
452-
LargeBinary | LargeUtf8 => 256,
453-
FixedSizeBinary(bytes) => (*bytes * 8) as usize,
454-
// primitive types
455-
Int8
456-
| Int16
457-
| Int32
458-
| Int64
459-
| UInt8
460-
| UInt16
461-
| UInt32
462-
| UInt64
463-
| Float16
464-
| Float32
465-
| Float64
466-
| Timestamp(_, _)
467-
| Date32
468-
| Date64
469-
| Time32(_)
470-
| Time64(_)
471-
| Duration(_)
472-
| Interval(_)
473-
| Dictionary(_, _)
474-
| Decimal128(_, _)
475-
| Decimal256(_, _) => data_type.primitive_width().unwrap_or(1) * 8,
476-
// nested types
477-
List(f) => nested_size(f.data_type(), nesting),
478-
FixedSizeList(_, s) => (s * 8) as usize,
479-
LargeList(f) => nested_size(f.data_type(), nesting),
480-
Struct(fields) => fields
481-
.iter()
482-
.map(|f| nested_size(f.data_type(), nesting))
483-
.sum(),
484-
Union(fields, _) => fields
485-
.iter()
486-
.map(|(_, f)| nested_size(f.data_type(), nesting))
487-
.sum(),
488-
Map(field, _) => nested_size(field.data_type(), nesting),
489-
RunEndEncoded(run_ends, values) => {
490-
nested_size(run_ends.data_type(), nesting)
491-
+ nested_size(values.data_type(), nesting)
492-
}
493-
}
494-
}
495-
496-
/// Find a field with a presumable small memory footprint based on its data type's memory size
497-
/// and the level of nesting.
498-
fn find_small_field(fields: &[DFField]) -> Option<DFField> {
499-
fields
500-
.iter()
501-
.map(|f| {
502-
let nesting = &mut 0;
503-
let size = nested_size(f.data_type(), nesting);
504-
(*nesting, size)
505-
})
506-
.enumerate()
507-
.min_by(|(_, (nesting_a, size_a)), (_, (nesting_b, size_b))| {
508-
nesting_a.cmp(nesting_b).then(size_a.cmp(size_b))
509-
})
510-
.map(|(i, _)| fields[i].clone())
511-
}
512-
513416
/// Get the projection exprs from columns in the order of the schema
514417
fn get_expr(columns: &HashSet<Column>, schema: &DFSchemaRef) -> Result<Vec<Expr>> {
515418
let expr = schema
@@ -640,7 +543,7 @@ mod tests {
640543
use crate::optimizer::Optimizer;
641544
use crate::test::*;
642545
use crate::OptimizerContext;
643-
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
546+
use arrow::datatypes::{DataType, Field, Schema};
644547
use datafusion_common::DFSchema;
645548
use datafusion_expr::builder::table_scan_with_filters;
646549
use datafusion_expr::expr;
@@ -1232,73 +1135,4 @@ mod tests {
12321135
.unwrap_or(optimized_plan);
12331136
Ok(optimized_plan)
12341137
}
1235-
1236-
#[test]
1237-
fn test_nested_size() {
1238-
use DataType::*;
1239-
let nesting = &mut 0;
1240-
assert_eq!(nested_size(&Null, nesting), 0);
1241-
assert_eq!(*nesting, 0);
1242-
assert_eq!(nested_size(&Boolean, nesting), 1);
1243-
assert_eq!(*nesting, 0);
1244-
assert_eq!(nested_size(&UInt8, nesting), 8);
1245-
assert_eq!(*nesting, 0);
1246-
assert_eq!(nested_size(&Int64, nesting), 64);
1247-
assert_eq!(*nesting, 0);
1248-
assert_eq!(nested_size(&Decimal256(5, 2), nesting), 256);
1249-
assert_eq!(*nesting, 0);
1250-
assert_eq!(
1251-
nested_size(&List(Arc::new(Field::new("A", Int64, true))), nesting),
1252-
64
1253-
);
1254-
assert_eq!(*nesting, 1);
1255-
*nesting = 0;
1256-
assert_eq!(
1257-
nested_size(
1258-
&List(Arc::new(Field::new(
1259-
"A",
1260-
List(Arc::new(Field::new("AA", Int64, true))),
1261-
true
1262-
))),
1263-
nesting
1264-
),
1265-
64
1266-
);
1267-
assert_eq!(*nesting, 2);
1268-
}
1269-
1270-
#[test]
1271-
fn test_find_small_field() {
1272-
use DataType::*;
1273-
let int32 = DFField::from(Field::new("a", Int32, false));
1274-
let bin = DFField::from(Field::new("b", Binary, false));
1275-
let list_i64 = DFField::from(Field::new(
1276-
"c",
1277-
List(Arc::new(Field::new("c_1", Int64, true))),
1278-
false,
1279-
));
1280-
let time_s = DFField::from(Field::new("d", Time32(TimeUnit::Second), false));
1281-
1282-
assert_eq!(
1283-
find_small_field(&[
1284-
int32.clone(),
1285-
bin.clone(),
1286-
list_i64.clone(),
1287-
time_s.clone()
1288-
]),
1289-
Some(int32.clone())
1290-
);
1291-
assert_eq!(
1292-
find_small_field(&[bin.clone(), list_i64.clone(), time_s.clone()]),
1293-
Some(time_s.clone())
1294-
);
1295-
assert_eq!(
1296-
find_small_field(&[time_s.clone(), int32.clone()]),
1297-
Some(time_s.clone())
1298-
);
1299-
assert_eq!(
1300-
find_small_field(&[bin.clone(), list_i64.clone()]),
1301-
Some(bin.clone())
1302-
);
1303-
}
13041138
}

‎datafusion/physical-plan/src/joins/cross_join.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ use crate::{
3434

3535
use arrow::datatypes::{Fields, Schema, SchemaRef};
3636
use arrow::record_batch::RecordBatch;
37+
use arrow_array::RecordBatchOptions;
3738
use datafusion_common::stats::Precision;
3839
use datafusion_common::{plan_err, DataFusionError, Result, ScalarValue};
3940
use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
@@ -347,13 +348,14 @@ fn build_batch(
347348
})
348349
.collect::<Result<Vec<_>>>()?;
349350

350-
RecordBatch::try_new(
351+
RecordBatch::try_new_with_options(
351352
Arc::new(schema.clone()),
352353
arrays
353354
.iter()
354355
.chain(batch.columns().iter())
355356
.cloned()
356357
.collect(),
358+
&RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
357359
)
358360
.map_err(Into::into)
359361
}

‎datafusion/sqllogictest/test_files/avro.slt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -253,10 +253,10 @@ EXPLAIN SELECT count(*) from alltypes_plain
253253
----
254254
logical_plan
255255
Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1)) AS COUNT(*)]]
256-
--TableScan: alltypes_plain projection=[bool_col]
256+
--TableScan: alltypes_plain projection=[]
257257
physical_plan
258258
AggregateExec: mode=Final, gby=[], aggr=[COUNT(*)]
259259
--CoalescePartitionsExec
260260
----AggregateExec: mode=Partial, gby=[], aggr=[COUNT(*)]
261261
------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
262-
--------AvroExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/avro/alltypes_plain.avro]]}, projection=[bool_col]
262+
--------AvroExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/avro/alltypes_plain.avro]]}

‎datafusion/sqllogictest/test_files/json.slt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,13 @@ EXPLAIN SELECT count(*) from json_test
5050
----
5151
logical_plan
5252
Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1)) AS COUNT(*)]]
53-
--TableScan: json_test projection=[c]
53+
--TableScan: json_test projection=[]
5454
physical_plan
5555
AggregateExec: mode=Final, gby=[], aggr=[COUNT(*)]
5656
--CoalescePartitionsExec
5757
----AggregateExec: mode=Partial, gby=[], aggr=[COUNT(*)]
5858
------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
59-
--------JsonExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/2.json]]}, projection=[c]
59+
--------JsonExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/2.json]]}
6060

6161
query ?
6262
SELECT mycol FROM single_nan

‎datafusion/sqllogictest/test_files/subquery.slt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -695,7 +695,7 @@ logical_plan
695695
Projection: __scalar_sq_1.COUNT(*) AS b
696696
--SubqueryAlias: __scalar_sq_1
697697
----Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1)) AS COUNT(*)]]
698-
------TableScan: t1 projection=[t1_id]
698+
------TableScan: t1 projection=[]
699699

700700
#simple_uncorrelated_scalar_subquery2
701701
query TT
@@ -706,10 +706,10 @@ Projection: __scalar_sq_1.COUNT(*) AS b, __scalar_sq_2.COUNT(Int64(1)) AS COUNT(
706706
--Left Join:
707707
----SubqueryAlias: __scalar_sq_1
708708
------Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1)) AS COUNT(*)]]
709-
--------TableScan: t1 projection=[t1_id]
709+
--------TableScan: t1 projection=[]
710710
----SubqueryAlias: __scalar_sq_2
711711
------Aggregate: groupBy=[[]], aggr=[[COUNT(Int64(1))]]
712-
--------TableScan: t2 projection=[t2_id]
712+
--------TableScan: t2 projection=[]
713713

714714
query II
715715
select (select count(*) from t1) as b, (select count(1) from t2)

0 commit comments

Comments
 (0)