Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,8 @@ impl CustomExec {
schema: SchemaRef,
db: CustomDataSource,
) -> Self {
let projected_schema = project_schema(&schema, projections).unwrap();
let projected_schema =
project_schema(&schema, projections.map(|v| v.as_ref())).unwrap();
let cache = Self::compute_properties(projected_schema.clone());
Self {
db,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/catalog-listing/src/table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -522,7 +522,7 @@ impl TableProvider for ListingTable {

// if no files need to be read, return an `EmptyExec`
if partitioned_file_lists.is_empty() {
let projected_schema = project_schema(&self.schema(), projection.as_ref())?;
let projected_schema = project_schema(&self.schema(), projection.as_deref())?;
return Ok(ScanResult::new(Arc::new(EmptyExec::new(projected_schema))));
}

Expand Down
16 changes: 8 additions & 8 deletions datafusion/common/src/stats.rs
Original file line number Diff line number Diff line change
Expand Up @@ -391,8 +391,8 @@ impl Statistics {
/// For example, if we had statistics for columns `{"a", "b", "c"}`,
/// projecting to `vec![2, 1]` would return statistics for columns `{"c",
/// "b"}`.
pub fn project(mut self, projection: Option<&Vec<usize>>) -> Self {
let Some(projection) = projection else {
pub fn project(mut self, projection: Option<&[usize]>) -> Self {
let Some(projection) = projection.map(AsRef::as_ref) else {
return self;
};

Expand All @@ -410,7 +410,7 @@ impl Statistics {
.map(Slot::Present)
.collect();

for idx in projection {
for idx in projection.iter() {
let next_idx = self.column_statistics.len();
let slot = std::mem::replace(
columns.get_mut(*idx).expect("projection out of bounds"),
Expand Down Expand Up @@ -1066,29 +1066,29 @@ mod tests {

#[test]
fn test_project_none() {
let projection = None;
let stats = make_stats(vec![10, 20, 30]).project(projection.as_ref());
let projection: Option<&[usize]> = None;
let stats = make_stats(vec![10, 20, 30]).project(projection);
assert_eq!(stats, make_stats(vec![10, 20, 30]));
}

#[test]
fn test_project_empty() {
let projection = Some(vec![]);
let stats = make_stats(vec![10, 20, 30]).project(projection.as_ref());
let stats = make_stats(vec![10, 20, 30]).project(projection.as_deref());
assert_eq!(stats, make_stats(vec![]));
}

#[test]
fn test_project_swap() {
let projection = Some(vec![2, 1]);
let stats = make_stats(vec![10, 20, 30]).project(projection.as_ref());
let stats = make_stats(vec![10, 20, 30]).project(projection.as_deref());
assert_eq!(stats, make_stats(vec![30, 20]));
}

#[test]
fn test_project_repeated() {
let projection = Some(vec![1, 2, 1, 1, 0, 2]);
let stats = make_stats(vec![10, 20, 30]).project(projection.as_ref());
let stats = make_stats(vec![10, 20, 30]).project(projection.as_deref());
assert_eq!(stats, make_stats(vec![20, 30, 20, 20, 10, 30]));
}

Expand Down
6 changes: 3 additions & 3 deletions datafusion/common/src/utils/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ use std::thread::available_parallelism;
///
/// // Pick columns 'c' and 'b'
/// let projection = Some(vec![2, 1]);
/// let projected_schema = project_schema(&schema, projection.as_ref()).unwrap();
/// let projected_schema = project_schema(&schema, projection.as_deref()).unwrap();
///
/// let expected_schema = SchemaRef::new(Schema::new(vec![
/// Field::new("c", DataType::Utf8, true),
Expand All @@ -70,9 +70,9 @@ use std::thread::available_parallelism;
/// ```
pub fn project_schema(
schema: &SchemaRef,
projection: Option<&Vec<usize>>,
projection: Option<&[usize]>,
) -> Result<SchemaRef> {
let schema = match projection {
let schema = match projection.map(AsRef::as_ref) {
Some(columns) => Arc::new(schema.project(columns)?),
None => Arc::clone(schema),
};
Expand Down
3 changes: 2 additions & 1 deletion datafusion/core/src/datasource/empty.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,8 @@ impl TableProvider for EmptyTable {
_limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
// even though there is no data, projections apply
let projected_schema = project_schema(&self.schema, projection)?;
let projected_schema =
project_schema(&self.schema, projection.map(AsRef::as_ref))?;
Ok(Arc::new(
EmptyExec::new(projected_schema).with_partitions(self.partitions),
))
Expand Down
2 changes: 1 addition & 1 deletion datafusion/core/src/physical_planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -984,7 +984,7 @@ impl DefaultPhysicalPlanner {
// project the output columns excluding the async functions
// The async functions are always appended to the end of the schema.
.apply_projection(Some(
(0..input.schema().fields().len()).collect(),
(0..input.schema().fields().len()).collect::<Vec<_>>(),
))?
.with_batch_size(session_state.config().batch_size())
.build()?
Expand Down
2 changes: 1 addition & 1 deletion datafusion/core/tests/custom_sources_cases/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ impl CustomExecutionPlan {
fn new(projection: Option<Vec<usize>>) -> Self {
let schema = TEST_CUSTOM_SCHEMA_REF!();
let schema =
project_schema(&schema, projection.as_ref()).expect("projected schema");
project_schema(&schema, projection.as_deref()).expect("projected schema");
let cache = Self::compute_properties(schema);
Self { projection, cache }
}
Expand Down
2 changes: 1 addition & 1 deletion datafusion/core/tests/physical_optimizer/join_selection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -762,7 +762,7 @@ async fn test_hash_join_swap_on_joins_with_projections(
"ProjectionExec won't be added above if HashJoinExec contains embedded projection",
);

assert_eq!(swapped_join.projection, Some(vec![0_usize]));
assert_eq!(swapped_join.projection.as_ref().unwrap(), [0_usize]);
assert_eq!(swapped.schema().fields.len(), 1);
assert_eq!(swapped.schema().fields[0].name(), "small_col");
Ok(())
Expand Down
2 changes: 1 addition & 1 deletion datafusion/datasource/src/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ impl MemorySourceConfig {
schema: SchemaRef,
projection: Option<Vec<usize>>,
) -> Result<Self> {
let projected_schema = project_schema(&schema, projection.as_ref())?;
let projected_schema = project_schema(&schema, projection.as_deref())?;
Ok(Self {
partitions: partitions.to_vec(),
schema,
Expand Down
Loading