1use std::ops::Deref;
21use std::sync::Arc;
22
23use crate::PhysicalExpr;
24use crate::expressions::{CastExpr, Column, Literal};
25use crate::scalar_function::ScalarFunctionExpr;
26use crate::utils::collect_columns;
27
28use arrow::array::{RecordBatch, RecordBatchOptions};
29use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
30use datafusion_common::stats::{ColumnStatistics, Precision};
31use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
32use datafusion_common::{
33 Result, ScalarValue, Statistics, assert_or_internal_err, internal_datafusion_err,
34 plan_err,
35};
36
37use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet;
38use datafusion_physical_expr_common::metrics::ExpressionEvaluatorMetrics;
39use datafusion_physical_expr_common::physical_expr::fmt_sql;
40use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
41use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays_with_metrics;
42use indexmap::IndexMap;
43use itertools::Itertools;
44
45#[derive(Debug, Clone)]
56pub struct ProjectionExpr {
57 pub expr: Arc<dyn PhysicalExpr>,
59 pub alias: String,
61}
62
63impl PartialEq for ProjectionExpr {
64 fn eq(&self, other: &Self) -> bool {
65 let ProjectionExpr { expr, alias } = self;
66 expr.eq(&other.expr) && *alias == other.alias
67 }
68}
69
70impl Eq for ProjectionExpr {}
71
72impl AsRef<Arc<dyn PhysicalExpr>> for ProjectionExpr {
75 fn as_ref(&self) -> &Arc<dyn PhysicalExpr> {
76 &self.expr
77 }
78}
79
80impl std::fmt::Display for ProjectionExpr {
81 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82 if self.expr.to_string() == self.alias {
83 write!(f, "{}", self.alias)
84 } else {
85 write!(f, "{} AS {}", self.expr, self.alias)
86 }
87 }
88}
89
90impl ProjectionExpr {
91 pub fn new(expr: Arc<dyn PhysicalExpr>, alias: impl Into<String>) -> Self {
93 let alias = alias.into();
94 Self { expr, alias }
95 }
96
97 pub fn new_from_expression(
99 expr: Arc<dyn PhysicalExpr>,
100 schema: &Schema,
101 ) -> Result<Self> {
102 let field = expr.return_field(schema)?;
103 Ok(Self {
104 expr,
105 alias: field.name().to_string(),
106 })
107 }
108}
109
110impl From<(Arc<dyn PhysicalExpr>, String)> for ProjectionExpr {
111 fn from(value: (Arc<dyn PhysicalExpr>, String)) -> Self {
112 Self::new(value.0, value.1)
113 }
114}
115
116impl From<&(Arc<dyn PhysicalExpr>, String)> for ProjectionExpr {
117 fn from(value: &(Arc<dyn PhysicalExpr>, String)) -> Self {
118 Self::new(Arc::clone(&value.0), value.1.clone())
119 }
120}
121
122impl From<ProjectionExpr> for (Arc<dyn PhysicalExpr>, String) {
123 fn from(value: ProjectionExpr) -> Self {
124 (value.expr, value.alias)
125 }
126}
127
128#[derive(Debug, Clone, PartialEq, Eq)]
137pub struct ProjectionExprs {
138 exprs: Arc<[ProjectionExpr]>,
140}
141
142impl std::fmt::Display for ProjectionExprs {
143 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
144 let exprs: Vec<String> = self.exprs.iter().map(|e| e.to_string()).collect();
145 write!(f, "Projection[{}]", exprs.join(", "))
146 }
147}
148
149impl From<Vec<ProjectionExpr>> for ProjectionExprs {
150 fn from(value: Vec<ProjectionExpr>) -> Self {
151 Self {
152 exprs: value.into(),
153 }
154 }
155}
156
157impl From<&[ProjectionExpr]> for ProjectionExprs {
158 fn from(value: &[ProjectionExpr]) -> Self {
159 Self {
160 exprs: value.iter().cloned().collect(),
161 }
162 }
163}
164
165impl FromIterator<ProjectionExpr> for ProjectionExprs {
166 fn from_iter<T: IntoIterator<Item = ProjectionExpr>>(exprs: T) -> Self {
167 Self {
168 exprs: exprs.into_iter().collect(),
169 }
170 }
171}
172
173impl AsRef<[ProjectionExpr]> for ProjectionExprs {
174 fn as_ref(&self) -> &[ProjectionExpr] {
175 &self.exprs
176 }
177}
178
179impl ProjectionExprs {
180 pub fn new(exprs: impl IntoIterator<Item = ProjectionExpr>) -> Self {
182 Self {
183 exprs: exprs.into_iter().collect(),
184 }
185 }
186
187 pub fn from_expressions(exprs: impl Into<Arc<[ProjectionExpr]>>) -> Self {
189 Self {
190 exprs: exprs.into(),
191 }
192 }
193
194 pub fn from_indices(indices: &[usize], schema: &Schema) -> Self {
238 let projection_exprs = indices.iter().map(|&i| {
239 let field = schema.field(i);
240 ProjectionExpr {
241 expr: Arc::new(Column::new(field.name(), i)),
242 alias: field.name().clone(),
243 }
244 });
245
246 Self::from_iter(projection_exprs)
247 }
248
249 pub fn iter(&self) -> impl Iterator<Item = &ProjectionExpr> {
251 self.exprs.iter()
252 }
253
254 pub fn projection_mapping(
256 &self,
257 input_schema: &SchemaRef,
258 ) -> Result<ProjectionMapping> {
259 ProjectionMapping::try_new(
260 self.exprs
261 .iter()
262 .map(|p| (Arc::clone(&p.expr), p.alias.clone())),
263 input_schema,
264 )
265 }
266
267 pub fn expr_iter(&self) -> impl Iterator<Item = Arc<dyn PhysicalExpr>> + '_ {
269 self.exprs.iter().map(|e| Arc::clone(&e.expr))
270 }
271
272 pub fn try_map_exprs<F>(self, mut f: F) -> Result<Self>
301 where
302 F: FnMut(Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>>,
303 {
304 let exprs = self
305 .exprs
306 .iter()
307 .cloned()
308 .map(|mut proj| {
309 proj.expr = f(proj.expr)?;
310 Ok(proj)
311 })
312 .collect::<Result<Arc<_>>>()?;
313 Ok(Self::from_expressions(exprs))
314 }
315
316 pub fn try_merge(&self, other: &ProjectionExprs) -> Result<ProjectionExprs> {
382 let mut new_exprs = Vec::with_capacity(other.exprs.len());
383 for proj_expr in other.exprs.iter() {
384 new_exprs.push(ProjectionExpr {
385 expr: self.unproject_expr(&proj_expr.expr)?,
386 alias: proj_expr.alias.clone(),
387 });
388 }
389 Ok(ProjectionExprs::new(new_exprs))
390 }
391
392 pub fn column_indices(&self) -> Vec<usize> {
397 self.exprs
398 .iter()
399 .flat_map(|e| collect_columns(&e.expr).into_iter().map(|col| col.index()))
400 .sorted_unstable()
401 .dedup()
402 .collect_vec()
403 }
404
405 #[deprecated(
437 since = "52.0.0",
438 note = "Use column_indices() instead. This method will be removed in 58.0.0 or 6 months after 52.0.0 is released, whichever comes first."
439 )]
440 pub fn ordered_column_indices(&self) -> Vec<usize> {
441 self.exprs
442 .iter()
443 .map(|e| {
444 e.expr
445 .downcast_ref::<Column>()
446 .expect("Expected column reference in projection")
447 .index()
448 })
449 .collect()
450 }
451
452 pub fn project_schema(&self, input_schema: &Schema) -> Result<Schema> {
464 let fields: Result<Vec<Field>> = self
465 .exprs
466 .iter()
467 .map(|proj_expr| {
468 let metadata = proj_expr
469 .expr
470 .return_field(input_schema)?
471 .metadata()
472 .clone();
473
474 let field = Field::new(
475 &proj_expr.alias,
476 proj_expr.expr.data_type(input_schema)?,
477 proj_expr.expr.nullable(input_schema)?,
478 )
479 .with_metadata(metadata);
480
481 Ok(field)
482 })
483 .collect();
484
485 Ok(Schema::new_with_metadata(
486 fields?,
487 input_schema.metadata().clone(),
488 ))
489 }
490
491 pub fn unproject_expr(
501 &self,
502 expr: &Arc<dyn PhysicalExpr>,
503 ) -> Result<Arc<dyn PhysicalExpr>> {
504 update_expr(expr, &self.exprs, true)?.ok_or_else(|| {
505 internal_datafusion_err!(
506 "Failed to unproject an expression {} with ProjectionExprs {}",
507 expr,
508 self.exprs.iter().map(|e| format!("{e}")).join(", ")
509 )
510 })
511 }
512
513 pub fn project_expr(
521 &self,
522 expr: &Arc<dyn PhysicalExpr>,
523 ) -> Result<Arc<dyn PhysicalExpr>> {
524 update_expr(expr, &self.exprs, false)?.ok_or_else(|| {
525 internal_datafusion_err!(
526 "Failed to project an expression {} with ProjectionExprs {}",
527 expr,
528 self.exprs.iter().map(|e| format!("{e}")).join(", ")
529 )
530 })
531 }
532
533 pub fn make_projector(&self, input_schema: &Schema) -> Result<Projector> {
542 let output_schema = Arc::new(self.project_schema(input_schema)?);
543 Ok(Projector {
544 projection: self.clone(),
545 output_schema,
546 expression_metrics: None,
547 })
548 }
549
550 pub fn make_projector_with_schema_metadata(
562 &self,
563 input_schema: &Schema,
564 projected_schema: &Schema,
565 ) -> Result<Projector> {
566 let output_schema = self.project_schema(input_schema)?;
567 if output_schema.fields().len() != projected_schema.fields().len() {
568 return Err(internal_datafusion_err!(
569 "Projection has {} output fields but metadata schema has {} fields",
570 output_schema.fields().len(),
571 projected_schema.fields().len()
572 ));
573 }
574
575 let fields = output_schema
576 .fields()
577 .iter()
578 .zip(projected_schema.fields())
579 .map(|(field, projected_field)| {
580 Arc::new(
581 field
582 .as_ref()
583 .clone()
584 .with_metadata(projected_field.metadata().clone()),
585 )
586 })
587 .collect::<Vec<_>>();
588 let output_schema = Arc::new(Schema::new_with_metadata(
589 fields,
590 projected_schema.metadata().clone(),
591 ));
592
593 Ok(Projector {
594 projection: self.clone(),
595 output_schema,
596 expression_metrics: None,
597 })
598 }
599
600 pub fn create_expression_metrics(
601 &self,
602 metrics: &ExecutionPlanMetricsSet,
603 partition: usize,
604 ) -> ExpressionEvaluatorMetrics {
605 let labels: Vec<String> = self
606 .exprs
607 .iter()
608 .map(|proj_expr| {
609 let expr_sql = fmt_sql(proj_expr.expr.as_ref()).to_string();
610 if proj_expr.expr.to_string() == proj_expr.alias {
611 expr_sql
612 } else {
613 format!("{expr_sql} AS {}", proj_expr.alias)
614 }
615 })
616 .collect();
617 ExpressionEvaluatorMetrics::new(metrics, partition, labels)
618 }
619
620 pub fn project_statistics(
713 &self,
714 mut stats: Statistics,
715 output_schema: &Schema,
716 ) -> Result<Statistics> {
717 let mut column_statistics = Vec::with_capacity(self.exprs.len());
718
719 for proj_expr in self.exprs.iter() {
720 let expr = &proj_expr.expr;
721 let col_stats = if let Some(col) = expr.downcast_ref::<Column>() {
722 column_statistics_at(&stats.column_statistics, col.index())
723 } else if let Some(literal) = expr.downcast_ref::<Literal>() {
724 let data_type = expr.data_type(output_schema)?;
726
727 if literal.value().is_null() {
728 let null_count = match stats.num_rows {
729 Precision::Exact(num_rows) => Precision::Exact(num_rows),
730 _ => Precision::Absent,
731 };
732
733 ColumnStatistics {
734 min_value: Precision::Exact(literal.value().clone()),
735 max_value: Precision::Exact(literal.value().clone()),
736 distinct_count: Precision::Exact(1),
737 null_count,
738 sum_value: Precision::Exact(literal.value().clone()),
739 byte_size: Precision::Exact(0),
740 }
741 } else {
742 let value = literal.value();
743 let distinct_count = Precision::Exact(1);
744 let null_count = Precision::Exact(0);
745
746 let byte_size = if let Some(byte_width) = data_type.primitive_width()
747 {
748 stats.num_rows.multiply(&Precision::Exact(byte_width))
749 } else {
750 Precision::Absent
752 };
753
754 let widened_sum = Precision::Exact(value.clone()).cast_to_sum_type();
755 let sum_value = widened_sum
756 .get_value()
757 .and_then(|sum| {
758 Precision::<ScalarValue>::from(stats.num_rows)
759 .cast_to(&sum.data_type())
760 .ok()
761 })
762 .map(|row_count| widened_sum.multiply(&row_count))
763 .unwrap_or(Precision::Absent);
764
765 ColumnStatistics {
766 min_value: Precision::Exact(value.clone()),
767 max_value: Precision::Exact(value.clone()),
768 distinct_count,
769 null_count,
770 sum_value,
771 byte_size,
772 }
773 }
774 } else {
775 project_column_statistics_through_expr(
776 expr.as_ref(),
777 &stats.column_statistics,
778 )
779 };
780 column_statistics.push(col_stats);
781 }
782 stats.calculate_total_byte_size(output_schema);
783 stats.column_statistics = column_statistics;
784 Ok(stats)
785 }
786
787 pub fn projected_column_position(&self, column: &Column) -> Option<usize> {
834 self.iter().position(|expr| {
835 expr.expr
836 .downcast_ref::<Column>()
837 .is_some_and(|projected| projected == column)
838 })
839 }
840}
841
842fn project_column_statistics_through_expr(
847 expr: &dyn PhysicalExpr,
848 column_stats: &[ColumnStatistics],
849) -> ColumnStatistics {
850 if let Some(col) = expr.downcast_ref::<Column>() {
851 return column_statistics_at(column_stats, col.index());
852 }
853 let Some(cast_expr) = expr.downcast_ref::<CastExpr>() else {
854 return ColumnStatistics::new_unknown();
855 };
856 let inner_stats =
857 project_column_statistics_through_expr(cast_expr.expr.as_ref(), column_stats);
858 let target_type = cast_expr.cast_type();
859
860 let already_target_type = matches!(
867 (inner_stats.min_value.get_value(), inner_stats.max_value.get_value()),
868 (Some(min), Some(max))
869 if min.data_type() == *target_type && max.data_type() == *target_type
870 );
871 if already_target_type {
872 return inner_stats;
873 }
874
875 ColumnStatistics {
876 min_value: inner_stats
877 .min_value
878 .cast_to(target_type)
879 .unwrap_or(Precision::Absent),
880 max_value: inner_stats
881 .max_value
882 .cast_to(target_type)
883 .unwrap_or(Precision::Absent),
884 null_count: inner_stats.null_count,
885 distinct_count: inner_stats.distinct_count,
886 sum_value: Precision::Absent,
887 byte_size: Precision::Absent,
888 }
889}
890
891fn column_statistics_at(
892 column_stats: &[ColumnStatistics],
893 index: usize,
894) -> ColumnStatistics {
895 column_stats
896 .get(index)
897 .cloned()
898 .unwrap_or_else(ColumnStatistics::new_unknown)
899}
900
901impl<'a> IntoIterator for &'a ProjectionExprs {
902 type Item = &'a ProjectionExpr;
903 type IntoIter = std::slice::Iter<'a, ProjectionExpr>;
904
905 fn into_iter(self) -> Self::IntoIter {
906 self.exprs.iter()
907 }
908}
909
910#[derive(Clone, Debug)]
919pub struct Projector {
920 projection: ProjectionExprs,
921 output_schema: SchemaRef,
922 expression_metrics: Option<ExpressionEvaluatorMetrics>,
924}
925
926impl Projector {
927 pub fn with_metrics(
932 &self,
933 metrics: &ExecutionPlanMetricsSet,
934 partition: usize,
935 ) -> Self {
936 let expr_metrics = self
937 .projection
938 .create_expression_metrics(metrics, partition);
939 Self {
940 expression_metrics: Some(expr_metrics),
941 projection: self.projection.clone(),
942 output_schema: Arc::clone(&self.output_schema),
943 }
944 }
945
946 pub fn project_batch(&self, batch: &RecordBatch) -> Result<RecordBatch> {
953 let arrays = evaluate_expressions_to_arrays_with_metrics(
954 self.projection.exprs.iter().map(|p| &p.expr),
955 batch,
956 self.expression_metrics.as_ref(),
957 )?;
958
959 if arrays.is_empty() {
960 let options =
961 RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
962 RecordBatch::try_new_with_options(
963 Arc::clone(&self.output_schema),
964 arrays,
965 &options,
966 )
967 .map_err(Into::into)
968 } else {
969 RecordBatch::try_new(Arc::clone(&self.output_schema), arrays)
970 .map_err(Into::into)
971 }
972 }
973
974 pub fn output_schema(&self) -> &SchemaRef {
975 &self.output_schema
976 }
977
978 pub fn projection(&self) -> &ProjectionExprs {
979 &self.projection
980 }
981}
982
983pub type ProjectionRef = Arc<[usize]>;
988
989pub fn combine_projections(
1005 p1: Option<&ProjectionRef>,
1006 p2: Option<&ProjectionRef>,
1007) -> Result<Option<ProjectionRef>> {
1008 let Some(p1) = p1 else {
1009 return Ok(None);
1010 };
1011 let Some(p2) = p2 else {
1012 return Ok(Some(Arc::clone(p1)));
1013 };
1014
1015 Ok(Some(
1016 p1.iter()
1017 .map(|i| {
1018 let idx = *i;
1019 assert_or_internal_err!(
1020 idx < p2.len(),
1021 "unable to apply projection: index {} is greater than new projection len {}",
1022 idx,
1023 p2.len(),
1024 );
1025 Ok(p2[*i])
1026 })
1027 .collect::<Result<Arc<[usize]>>>()?,
1028 ))
1029}
1030
1031pub fn update_expr(
1077 expr: &Arc<dyn PhysicalExpr>,
1078 projected_exprs: &[ProjectionExpr],
1079 unproject: bool,
1080) -> Result<Option<Arc<dyn PhysicalExpr>>> {
1081 #[derive(Debug, PartialEq)]
1082 enum RewriteState {
1083 Unchanged,
1085 RewrittenValid,
1087 RewrittenInvalid,
1090 }
1091
1092 let mut state = RewriteState::Unchanged;
1093
1094 let new_expr = Arc::clone(expr)
1095 .transform_up(|expr| {
1096 if state == RewriteState::RewrittenInvalid {
1097 return Ok(Transformed::no(expr));
1098 }
1099
1100 let Some(column) = expr.downcast_ref::<Column>() else {
1101 return Ok(Transformed::no(expr));
1102 };
1103 if unproject {
1104 let projected_expr = projected_exprs.get(column.index()).ok_or_else(|| {
1105 internal_datafusion_err!(
1106 "Column index {} out of bounds for projected expressions of length {}",
1107 column.index(),
1108 projected_exprs.len()
1109 )
1110 })?;
1111 if let Some(projected_col) =
1116 projected_expr.expr.downcast_ref::<Column>()
1117 && projected_col == column
1118 {
1119 return Ok(Transformed::no(expr));
1120 }
1121 state = RewriteState::RewrittenValid;
1122 Ok(Transformed::yes(Arc::clone(&projected_expr.expr)))
1123 } else {
1124 state = RewriteState::RewrittenInvalid;
1126 projected_exprs
1128 .iter()
1129 .enumerate()
1130 .find_map(|(index, proj_expr)| {
1131 proj_expr.expr.downcast_ref::<Column>().and_then(
1132 |projected_column| {
1133 (column.name().eq(projected_column.name())
1134 && column.index() == projected_column.index())
1135 .then(|| {
1136 state = RewriteState::RewrittenValid;
1137 Arc::new(Column::new(&proj_expr.alias, index)) as _
1138 })
1139 },
1140 )
1141 })
1142 .map_or_else(
1143 || Ok(Transformed::no(expr)),
1144 |c| Ok(Transformed::yes(c)),
1145 )
1146 }
1147 })
1148 .data()?;
1149
1150 match state {
1151 RewriteState::RewrittenInvalid => Ok(None),
1152 RewriteState::Unchanged | RewriteState::RewrittenValid => Ok(Some(new_expr)),
1156 }
1157}
1158
1159#[derive(Clone, Debug, Default)]
1162pub struct ProjectionTargets {
1163 exprs_indices: Vec<(Arc<dyn PhysicalExpr>, usize)>,
1167}
1168
1169impl ProjectionTargets {
1170 pub fn first(&self) -> &(Arc<dyn PhysicalExpr>, usize) {
1172 self.exprs_indices.first().unwrap()
1174 }
1175
1176 pub fn push(&mut self, target: (Arc<dyn PhysicalExpr>, usize)) {
1178 self.exprs_indices.push(target);
1179 }
1180}
1181
1182impl Deref for ProjectionTargets {
1183 type Target = [(Arc<dyn PhysicalExpr>, usize)];
1184
1185 fn deref(&self) -> &Self::Target {
1186 &self.exprs_indices
1187 }
1188}
1189
1190impl From<Vec<(Arc<dyn PhysicalExpr>, usize)>> for ProjectionTargets {
1191 fn from(exprs_indices: Vec<(Arc<dyn PhysicalExpr>, usize)>) -> Self {
1192 Self { exprs_indices }
1193 }
1194}
1195
1196#[derive(Clone, Debug)]
1199pub struct ProjectionMapping {
1200 map: IndexMap<Arc<dyn PhysicalExpr>, ProjectionTargets>,
1203}
1204
1205impl ProjectionMapping {
1206 pub fn try_new(
1220 expr: impl IntoIterator<Item = (Arc<dyn PhysicalExpr>, String)>,
1221 input_schema: &SchemaRef,
1222 ) -> Result<Self> {
1223 let mut map = IndexMap::<_, ProjectionTargets>::new();
1225 for (expr_idx, (expr, name)) in expr.into_iter().enumerate() {
1226 let target_expr = Arc::new(Column::new(&name, expr_idx)) as _;
1227 let source_expr = expr.transform_down(|e| match e.downcast_ref::<Column>() {
1228 Some(col) => {
1229 let idx = col.index();
1234 let matching_field = input_schema.field(idx);
1235 let matching_name = matching_field.name();
1236 assert_or_internal_err!(
1237 col.name() == matching_name,
1238 "Input field name {matching_name} does not match with the projection expression {}",
1239 col.name()
1240 );
1241 let matching_column = Column::new(matching_name, idx);
1242 Ok(Transformed::yes(Arc::new(matching_column)))
1243 }
1244 None => Ok(Transformed::no(e)),
1245 })
1246 .data()?;
1247 map.entry(Arc::clone(&source_expr))
1248 .or_default()
1249 .push((Arc::clone(&target_expr), expr_idx));
1250
1251 if let Some(func_expr) = source_expr.downcast_ref::<ScalarFunctionExpr>() {
1259 let literal_args: Vec<Option<ScalarValue>> = func_expr
1260 .args()
1261 .iter()
1262 .map(|arg| arg.downcast_ref::<Literal>().map(|l| l.value().clone()))
1263 .collect();
1264
1265 if let Some(field_mapping) =
1266 func_expr.fun().struct_field_mapping(&literal_args)
1267 && let DataType::Struct(struct_fields) = func_expr.return_type()
1268 {
1269 for (accessor_args, source_arg_idx) in &field_mapping.fields {
1270 let value_expr = Arc::clone(&func_expr.args()[*source_arg_idx]);
1271
1272 let mut accessor_fn_args: Vec<Arc<dyn PhysicalExpr>> =
1274 vec![Arc::clone(&target_expr)];
1275 accessor_fn_args.extend(accessor_args.iter().map(|sv| {
1276 Arc::new(Literal::new(sv.clone())) as Arc<dyn PhysicalExpr>
1277 }));
1278
1279 let return_field = accessor_args
1281 .first()
1282 .and_then(|sv| sv.try_as_str().flatten())
1283 .and_then(|field_name| {
1284 struct_fields
1285 .iter()
1286 .find(|f| f.name() == field_name)
1287 .cloned()
1288 });
1289
1290 if let Some(return_field) = return_field {
1291 let field_access_expr = Arc::new(ScalarFunctionExpr::new(
1292 field_mapping.field_accessor.name(),
1293 Arc::clone(&field_mapping.field_accessor),
1294 accessor_fn_args,
1295 return_field,
1296 Arc::new(func_expr.config_options().clone()),
1297 ))
1298 as Arc<dyn PhysicalExpr>;
1299
1300 map.entry(value_expr)
1301 .or_default()
1302 .push((field_access_expr, expr_idx));
1303 }
1304 }
1305 }
1306 }
1307 }
1308 Ok(Self { map })
1309 }
1310
1311 pub fn from_indices(indices: &[usize], schema: &SchemaRef) -> Result<Self> {
1316 let projection_exprs = indices.iter().map(|index| {
1317 let field = schema.field(*index);
1318 let column = Arc::new(Column::new(field.name(), *index));
1319 (column as _, field.name().clone())
1320 });
1321 ProjectionMapping::try_new(projection_exprs, schema)
1322 }
1323}
1324
1325impl Deref for ProjectionMapping {
1326 type Target = IndexMap<Arc<dyn PhysicalExpr>, ProjectionTargets>;
1327
1328 fn deref(&self) -> &Self::Target {
1329 &self.map
1330 }
1331}
1332
1333impl FromIterator<(Arc<dyn PhysicalExpr>, ProjectionTargets)> for ProjectionMapping {
1334 fn from_iter<T: IntoIterator<Item = (Arc<dyn PhysicalExpr>, ProjectionTargets)>>(
1335 iter: T,
1336 ) -> Self {
1337 Self {
1338 map: IndexMap::from_iter(iter),
1339 }
1340 }
1341}
1342
1343pub fn project_orderings(
1356 orderings: &[LexOrdering],
1357 schema: &SchemaRef,
1358) -> Vec<LexOrdering> {
1359 let mut projected_orderings = vec![];
1360
1361 for ordering in orderings {
1362 projected_orderings.extend(project_ordering(ordering, schema));
1363 }
1364
1365 projected_orderings
1366}
1367
1368pub fn project_ordering(
1398 ordering: &LexOrdering,
1399 schema: &SchemaRef,
1400) -> Option<LexOrdering> {
1401 let mut projected_exprs = vec![];
1402 for PhysicalSortExpr { expr, options } in ordering.iter() {
1403 let transformed = Arc::clone(expr).transform_up(|expr| {
1404 let Some(col) = expr.downcast_ref::<Column>() else {
1405 return Ok(Transformed::no(expr));
1406 };
1407
1408 let name = col.name();
1409 if let Some((idx, _)) = schema.column_with_name(name) {
1410 Ok(Transformed::yes(Arc::new(Column::new(name, idx))))
1412 } else {
1413 plan_err!("")
1416 }
1417 });
1418
1419 match transformed {
1420 Ok(transformed) => {
1421 projected_exprs.push(PhysicalSortExpr::new(transformed.data, *options));
1422 }
1423 Err(_) => {
1424 break;
1427 }
1428 }
1429 }
1430
1431 LexOrdering::new(projected_exprs)
1432}
1433
1434#[cfg(test)]
1435pub(crate) mod tests {
1436 use std::collections::HashMap;
1437
1438 use super::*;
1439 use crate::equivalence::{EquivalenceProperties, convert_to_orderings};
1440 use crate::expressions::{BinaryExpr, CastExpr, col};
1441 use crate::utils::tests::TestScalarUDF;
1442 use crate::{PhysicalExprRef, ScalarFunctionExpr};
1443
1444 use arrow::compute::SortOptions;
1445 use arrow::datatypes::{DataType, TimeUnit};
1446 use datafusion_common::config::ConfigOptions;
1447 use datafusion_expr::{Operator, ScalarUDF};
1448 use insta::assert_snapshot;
1449
1450 pub(crate) fn output_schema(
1451 mapping: &ProjectionMapping,
1452 input_schema: &Arc<Schema>,
1453 ) -> Result<SchemaRef> {
1454 let mut fields = vec![];
1456 for (source, targets) in mapping.iter() {
1457 let data_type = source.data_type(input_schema)?;
1458 let nullable = source.nullable(input_schema)?;
1459 for (target, _) in targets.iter() {
1460 let Some(column) = target.downcast_ref::<Column>() else {
1463 continue;
1464 };
1465 fields.push(Field::new(column.name(), data_type.clone(), nullable));
1466 }
1467 }
1468
1469 let output_schema = Arc::new(Schema::new_with_metadata(
1470 fields,
1471 input_schema.metadata().clone(),
1472 ));
1473
1474 Ok(output_schema)
1475 }
1476
1477 #[test]
1478 fn project_orderings() -> Result<()> {
1479 let schema = Arc::new(Schema::new(vec![
1480 Field::new("a", DataType::Int32, true),
1481 Field::new("b", DataType::Int32, true),
1482 Field::new("c", DataType::Int32, true),
1483 Field::new("d", DataType::Int32, true),
1484 Field::new("e", DataType::Int32, true),
1485 Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
1486 ]));
1487 let col_a = &col("a", &schema)?;
1488 let col_b = &col("b", &schema)?;
1489 let col_c = &col("c", &schema)?;
1490 let col_d = &col("d", &schema)?;
1491 let col_e = &col("e", &schema)?;
1492 let col_ts = &col("ts", &schema)?;
1493 let a_plus_b = Arc::new(BinaryExpr::new(
1494 Arc::clone(col_a),
1495 Operator::Plus,
1496 Arc::clone(col_b),
1497 )) as Arc<dyn PhysicalExpr>;
1498 let b_plus_d = Arc::new(BinaryExpr::new(
1499 Arc::clone(col_b),
1500 Operator::Plus,
1501 Arc::clone(col_d),
1502 )) as Arc<dyn PhysicalExpr>;
1503 let b_plus_e = Arc::new(BinaryExpr::new(
1504 Arc::clone(col_b),
1505 Operator::Plus,
1506 Arc::clone(col_e),
1507 )) as Arc<dyn PhysicalExpr>;
1508 let c_plus_d = Arc::new(BinaryExpr::new(
1509 Arc::clone(col_c),
1510 Operator::Plus,
1511 Arc::clone(col_d),
1512 )) as Arc<dyn PhysicalExpr>;
1513
1514 let option_asc = SortOptions {
1515 descending: false,
1516 nulls_first: false,
1517 };
1518 let option_desc = SortOptions {
1519 descending: true,
1520 nulls_first: true,
1521 };
1522
1523 let test_cases = vec![
1524 (
1526 vec![
1528 vec![(col_b, option_asc)],
1530 ],
1531 vec![(col_b, "b_new".to_string()), (col_a, "a_new".to_string())],
1533 vec![
1535 vec![("b_new", option_asc)],
1537 ],
1538 ),
1539 (
1541 vec![
1543 ],
1545 vec![(col_c, "c_new".to_string()), (col_b, "b_new".to_string())],
1547 vec![
1549 ],
1551 ),
1552 (
1554 vec![
1556 vec![(col_ts, option_asc)],
1558 ],
1559 vec![
1561 (col_b, "b_new".to_string()),
1562 (col_a, "a_new".to_string()),
1563 (col_ts, "ts_new".to_string()),
1564 ],
1565 vec![
1567 vec![("ts_new", option_asc)],
1569 ],
1570 ),
1571 (
1573 vec![
1575 vec![(col_a, option_asc), (col_ts, option_asc)],
1577 vec![(col_b, option_asc), (col_ts, option_asc)],
1579 ],
1580 vec![
1582 (col_b, "b_new".to_string()),
1583 (col_a, "a_new".to_string()),
1584 (col_ts, "ts_new".to_string()),
1585 ],
1586 vec![
1588 vec![("a_new", option_asc), ("ts_new", option_asc)],
1590 vec![("b_new", option_asc), ("ts_new", option_asc)],
1592 ],
1593 ),
1594 (
1596 vec![
1598 vec![(&a_plus_b, option_asc)],
1600 ],
1601 vec![
1603 (col_b, "b_new".to_string()),
1604 (col_a, "a_new".to_string()),
1605 (&a_plus_b, "a+b".to_string()),
1606 ],
1607 vec![
1609 vec![("a+b", option_asc)],
1611 ],
1612 ),
1613 (
1615 vec![
1617 vec![(&a_plus_b, option_asc), (col_c, option_asc)],
1619 ],
1620 vec![
1622 (col_b, "b_new".to_string()),
1623 (col_a, "a_new".to_string()),
1624 (col_c, "c_new".to_string()),
1625 (&a_plus_b, "a+b".to_string()),
1626 ],
1627 vec![
1629 vec![("a+b", option_asc), ("c_new", option_asc)],
1631 ],
1632 ),
1633 (
1635 vec![
1636 vec![(col_a, option_asc), (col_b, option_asc)],
1638 vec![(col_a, option_asc), (col_d, option_asc)],
1640 ],
1641 vec![
1643 (col_b, "b_new".to_string()),
1644 (col_a, "a_new".to_string()),
1645 (col_d, "d_new".to_string()),
1646 (&b_plus_d, "b+d".to_string()),
1647 ],
1648 vec![
1650 vec![("a_new", option_asc), ("b_new", option_asc)],
1652 vec![("a_new", option_asc), ("d_new", option_asc)],
1654 ],
1655 ),
1656 (
1658 vec![
1660 vec![(&b_plus_d, option_asc)],
1662 ],
1663 vec![
1665 (col_b, "b_new".to_string()),
1666 (col_a, "a_new".to_string()),
1667 (col_d, "d_new".to_string()),
1668 (&b_plus_d, "b+d".to_string()),
1669 ],
1670 vec![
1672 vec![("b+d", option_asc)],
1674 ],
1675 ),
1676 (
1678 vec![
1680 vec![
1682 (col_a, option_asc),
1683 (col_d, option_asc),
1684 (col_b, option_asc),
1685 ],
1686 vec![(col_c, option_asc)],
1688 ],
1689 vec![
1691 (col_b, "b_new".to_string()),
1692 (col_a, "a_new".to_string()),
1693 (col_d, "d_new".to_string()),
1694 (col_c, "c_new".to_string()),
1695 ],
1696 vec![
1698 vec![
1700 ("a_new", option_asc),
1701 ("d_new", option_asc),
1702 ("b_new", option_asc),
1703 ],
1704 vec![("c_new", option_asc)],
1706 ],
1707 ),
1708 (
1710 vec![
1711 vec![
1713 (col_a, option_asc),
1714 (col_b, option_asc),
1715 (col_c, option_asc),
1716 ],
1717 vec![(col_a, option_asc), (col_d, option_asc)],
1719 ],
1720 vec![
1722 (col_b, "b_new".to_string()),
1723 (col_a, "a_new".to_string()),
1724 (col_c, "c_new".to_string()),
1725 (&c_plus_d, "c+d".to_string()),
1726 ],
1727 vec![
1729 vec![
1731 ("a_new", option_asc),
1732 ("b_new", option_asc),
1733 ("c_new", option_asc),
1734 ],
1735 ],
1736 ),
1737 (
1739 vec![
1741 vec![(col_a, option_asc), (col_b, option_asc)],
1743 vec![(col_a, option_asc), (col_d, option_asc)],
1745 ],
1746 vec![
1748 (col_b, "b_new".to_string()),
1749 (col_a, "a_new".to_string()),
1750 (&b_plus_d, "b+d".to_string()),
1751 ],
1752 vec![
1754 vec![("a_new", option_asc), ("b_new", option_asc)],
1756 ],
1757 ),
1758 (
1760 vec![
1762 vec![
1764 (col_a, option_asc),
1765 (col_b, option_asc),
1766 (col_c, option_asc),
1767 ],
1768 ],
1769 vec![(col_c, "c_new".to_string()), (col_a, "a_new".to_string())],
1771 vec![
1773 vec![("a_new", option_asc)],
1775 ],
1776 ),
1777 (
1779 vec![
1781 vec![
1783 (col_a, option_asc),
1784 (col_b, option_asc),
1785 (col_c, option_asc),
1786 ],
1787 vec![
1789 (col_a, option_asc),
1790 (&a_plus_b, option_asc),
1791 (col_c, option_asc),
1792 ],
1793 ],
1794 vec![
1796 (col_c, "c_new".to_string()),
1797 (col_b, "b_new".to_string()),
1798 (col_a, "a_new".to_string()),
1799 (&a_plus_b, "a+b".to_string()),
1800 ],
1801 vec![
1803 vec![
1805 ("a_new", option_asc),
1806 ("b_new", option_asc),
1807 ("c_new", option_asc),
1808 ],
1809 vec![
1811 ("a_new", option_asc),
1812 ("a+b", option_asc),
1813 ("c_new", option_asc),
1814 ],
1815 ],
1816 ),
1817 (
1819 vec![
1821 vec![(col_a, option_asc), (col_b, option_asc)],
1823 vec![(col_c, option_asc), (col_b, option_asc)],
1825 vec![(col_d, option_asc), (col_e, option_asc)],
1827 ],
1828 vec![
1830 (col_c, "c_new".to_string()),
1831 (col_d, "d_new".to_string()),
1832 (col_a, "a_new".to_string()),
1833 (&b_plus_e, "b+e".to_string()),
1834 ],
1835 vec![
1837 vec![("a_new", option_asc)],
1839 vec![("c_new", option_asc)],
1841 vec![("d_new", option_asc)],
1843 ],
1844 ),
1845 (
1847 vec![
1849 vec![
1851 (col_a, option_asc),
1852 (col_c, option_asc),
1853 (col_b, option_asc),
1854 ],
1855 ],
1856 vec![
1858 (col_c, "c_new".to_string()),
1859 (col_a, "a_new".to_string()),
1860 (&a_plus_b, "a+b".to_string()),
1861 ],
1862 vec![
1864 vec![("a_new", option_asc), ("c_new", option_asc)],
1866 ],
1867 ),
1868 (
1870 vec![
1872 vec![(col_a, option_asc), (col_b, option_asc)],
1874 vec![(col_c, option_asc), (col_b, option_desc)],
1876 vec![(col_e, option_asc)],
1878 ],
1879 vec![
1881 (col_c, "c_new".to_string()),
1882 (col_a, "a_new".to_string()),
1883 (col_b, "b_new".to_string()),
1884 (&b_plus_e, "b+e".to_string()),
1885 ],
1886 vec![
1888 vec![("a_new", option_asc), ("b_new", option_asc)],
1890 vec![("c_new", option_asc), ("b_new", option_desc)],
1892 ],
1893 ),
1894 ];
1895
1896 for (idx, (orderings, proj_exprs, expected)) in test_cases.into_iter().enumerate()
1897 {
1898 let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
1899
1900 let orderings = convert_to_orderings(&orderings);
1901 eq_properties.add_orderings(orderings);
1902
1903 let proj_exprs = proj_exprs
1904 .into_iter()
1905 .map(|(expr, name)| (Arc::clone(expr), name));
1906 let projection_mapping = ProjectionMapping::try_new(proj_exprs, &schema)?;
1907 let output_schema = output_schema(&projection_mapping, &schema)?;
1908
1909 let expected = expected
1910 .into_iter()
1911 .map(|ordering| {
1912 ordering
1913 .into_iter()
1914 .map(|(name, options)| {
1915 (col(name, &output_schema).unwrap(), options)
1916 })
1917 .collect::<Vec<_>>()
1918 })
1919 .collect::<Vec<_>>();
1920 let expected = convert_to_orderings(&expected);
1921
1922 let projected_eq = eq_properties.project(&projection_mapping, output_schema);
1923 let orderings = projected_eq.oeq_class();
1924
1925 let err_msg = format!(
1926 "test_idx: {idx:?}, actual: {orderings:?}, expected: {expected:?}, projection_mapping: {projection_mapping:?}"
1927 );
1928
1929 assert_eq!(orderings.len(), expected.len(), "{err_msg}");
1930 for expected_ordering in &expected {
1931 assert!(orderings.contains(expected_ordering), "{}", err_msg)
1932 }
1933 }
1934
1935 Ok(())
1936 }
1937
1938 #[test]
1939 fn project_orderings2() -> Result<()> {
1940 let schema = Arc::new(Schema::new(vec![
1941 Field::new("a", DataType::Int32, true),
1942 Field::new("b", DataType::Int32, true),
1943 Field::new("c", DataType::Int32, true),
1944 Field::new("d", DataType::Int32, true),
1945 Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
1946 ]));
1947 let col_a = &col("a", &schema)?;
1948 let col_b = &col("b", &schema)?;
1949 let col_c = &col("c", &schema)?;
1950 let col_ts = &col("ts", &schema)?;
1951 let a_plus_b = Arc::new(BinaryExpr::new(
1952 Arc::clone(col_a),
1953 Operator::Plus,
1954 Arc::clone(col_b),
1955 )) as Arc<dyn PhysicalExpr>;
1956
1957 let test_fun = Arc::new(ScalarUDF::new_from_impl(TestScalarUDF::new()));
1958
1959 let round_c = Arc::new(ScalarFunctionExpr::try_new(
1960 test_fun,
1961 vec![Arc::clone(col_c)],
1962 &schema,
1963 Arc::new(ConfigOptions::default()),
1964 )?) as PhysicalExprRef;
1965
1966 let option_asc = SortOptions {
1967 descending: false,
1968 nulls_first: false,
1969 };
1970
1971 let proj_exprs = vec![
1972 (col_b, "b_new".to_string()),
1973 (col_a, "a_new".to_string()),
1974 (col_c, "c_new".to_string()),
1975 (&round_c, "round_c_res".to_string()),
1976 ];
1977 let proj_exprs = proj_exprs
1978 .into_iter()
1979 .map(|(expr, name)| (Arc::clone(expr), name));
1980 let projection_mapping = ProjectionMapping::try_new(proj_exprs, &schema)?;
1981 let output_schema = output_schema(&projection_mapping, &schema)?;
1982
1983 let col_a_new = &col("a_new", &output_schema)?;
1984 let col_b_new = &col("b_new", &output_schema)?;
1985 let col_c_new = &col("c_new", &output_schema)?;
1986 let col_round_c_res = &col("round_c_res", &output_schema)?;
1987 let a_new_plus_b_new = Arc::new(BinaryExpr::new(
1988 Arc::clone(col_a_new),
1989 Operator::Plus,
1990 Arc::clone(col_b_new),
1991 )) as Arc<dyn PhysicalExpr>;
1992
1993 let test_cases = [
1994 (
1996 vec![
1998 vec![(col_a, option_asc)],
2000 ],
2001 vec![
2003 vec![(col_a_new, option_asc)],
2005 ],
2006 ),
2007 (
2009 vec![
2011 vec![(&a_plus_b, option_asc)],
2013 ],
2014 vec![
2016 vec![(&a_new_plus_b_new, option_asc)],
2018 ],
2019 ),
2020 (
2022 vec![
2024 vec![(col_a, option_asc), (col_ts, option_asc)],
2026 ],
2027 vec![
2029 vec![(col_a_new, option_asc)],
2031 ],
2032 ),
2033 (
2035 vec![
2037 vec![
2039 (col_a, option_asc),
2040 (col_ts, option_asc),
2041 (col_b, option_asc),
2042 ],
2043 ],
2044 vec![
2046 vec![(col_a_new, option_asc)],
2048 ],
2049 ),
2050 (
2052 vec![
2054 vec![(col_a, option_asc), (col_c, option_asc)],
2056 ],
2057 vec![
2059 vec![(col_a_new, option_asc), (col_round_c_res, option_asc)],
2061 vec![(col_a_new, option_asc), (col_c_new, option_asc)],
2063 ],
2064 ),
2065 (
2067 vec![
2069 vec![(col_c, option_asc), (col_b, option_asc)],
2071 ],
2072 vec![
2074 vec![(col_round_c_res, option_asc)],
2076 vec![(col_c_new, option_asc), (col_b_new, option_asc)],
2078 ],
2079 ),
2080 (
2082 vec![
2084 vec![(&a_plus_b, option_asc), (col_c, option_asc)],
2086 ],
2087 vec![
2089 vec![
2091 (&a_new_plus_b_new, option_asc),
2092 (col_round_c_res, option_asc),
2093 ],
2094 vec![(&a_new_plus_b_new, option_asc), (col_c_new, option_asc)],
2096 ],
2097 ),
2098 ];
2099
2100 for (idx, (orderings, expected)) in test_cases.iter().enumerate() {
2101 let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
2102
2103 let orderings = convert_to_orderings(orderings);
2104 eq_properties.add_orderings(orderings);
2105
2106 let expected = convert_to_orderings(expected);
2107
2108 let projected_eq =
2109 eq_properties.project(&projection_mapping, Arc::clone(&output_schema));
2110 let orderings = projected_eq.oeq_class();
2111
2112 let err_msg = format!(
2113 "test idx: {idx:?}, actual: {orderings:?}, expected: {expected:?}, projection_mapping: {projection_mapping:?}"
2114 );
2115
2116 assert_eq!(orderings.len(), expected.len(), "{err_msg}");
2117 for expected_ordering in &expected {
2118 assert!(orderings.contains(expected_ordering), "{}", err_msg)
2119 }
2120 }
2121 Ok(())
2122 }
2123
2124 #[test]
2125 fn project_orderings3() -> Result<()> {
2126 let schema = Arc::new(Schema::new(vec![
2127 Field::new("a", DataType::Int32, true),
2128 Field::new("b", DataType::Int32, true),
2129 Field::new("c", DataType::Int32, true),
2130 Field::new("d", DataType::Int32, true),
2131 Field::new("e", DataType::Int32, true),
2132 Field::new("f", DataType::Int32, true),
2133 ]));
2134 let col_a = &col("a", &schema)?;
2135 let col_b = &col("b", &schema)?;
2136 let col_c = &col("c", &schema)?;
2137 let col_d = &col("d", &schema)?;
2138 let col_e = &col("e", &schema)?;
2139 let col_f = &col("f", &schema)?;
2140 let a_plus_b = Arc::new(BinaryExpr::new(
2141 Arc::clone(col_a),
2142 Operator::Plus,
2143 Arc::clone(col_b),
2144 )) as Arc<dyn PhysicalExpr>;
2145
2146 let option_asc = SortOptions {
2147 descending: false,
2148 nulls_first: false,
2149 };
2150
2151 let proj_exprs = vec![
2152 (col_c, "c_new".to_string()),
2153 (col_d, "d_new".to_string()),
2154 (&a_plus_b, "a+b".to_string()),
2155 ];
2156 let proj_exprs = proj_exprs
2157 .into_iter()
2158 .map(|(expr, name)| (Arc::clone(expr), name));
2159 let projection_mapping = ProjectionMapping::try_new(proj_exprs, &schema)?;
2160 let output_schema = output_schema(&projection_mapping, &schema)?;
2161
2162 let col_c_new = &col("c_new", &output_schema)?;
2163 let col_d_new = &col("d_new", &output_schema)?;
2164
2165 let test_cases = vec![
2166 (
2168 vec![
2170 vec![(col_d, option_asc), (col_b, option_asc)],
2172 vec![(col_c, option_asc), (col_a, option_asc)],
2174 ],
2175 vec![],
2177 vec![
2179 vec![(col_c_new, option_asc)],
2181 vec![(col_d_new, option_asc)],
2183 ],
2184 ),
2185 (
2187 vec![
2189 vec![(col_d, option_asc), (col_b, option_asc)],
2191 vec![(col_c, option_asc), (col_e, option_asc)],
2193 ],
2194 vec![(col_e, col_a)],
2196 vec![
2198 vec![(col_c_new, option_asc)],
2200 vec![(col_d_new, option_asc)],
2202 ],
2203 ),
2204 (
2206 vec![
2208 vec![(col_d, option_asc), (col_b, option_asc)],
2210 vec![(col_c, option_asc), (col_e, option_asc)],
2212 ],
2213 vec![(col_a, col_f)],
2215 vec![
2217 vec![(col_d_new, option_asc)],
2219 vec![(col_c_new, option_asc)],
2221 ],
2222 ),
2223 ];
2224 for (orderings, equal_columns, expected) in test_cases {
2225 let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
2226 for (lhs, rhs) in equal_columns {
2227 eq_properties.add_equal_conditions(Arc::clone(lhs), Arc::clone(rhs))?;
2228 }
2229
2230 let orderings = convert_to_orderings(&orderings);
2231 eq_properties.add_orderings(orderings);
2232
2233 let expected = convert_to_orderings(&expected);
2234
2235 let projected_eq =
2236 eq_properties.project(&projection_mapping, Arc::clone(&output_schema));
2237 let orderings = projected_eq.oeq_class();
2238
2239 let err_msg = format!(
2240 "actual: {orderings:?}, expected: {expected:?}, projection_mapping: {projection_mapping:?}"
2241 );
2242
2243 assert_eq!(orderings.len(), expected.len(), "{err_msg}");
2244 for expected_ordering in &expected {
2245 assert!(orderings.contains(expected_ordering), "{}", err_msg)
2246 }
2247 }
2248
2249 Ok(())
2250 }
2251
2252 fn get_stats() -> Statistics {
2253 Statistics {
2254 num_rows: Precision::Exact(5),
2255 total_byte_size: Precision::Exact(23),
2256 column_statistics: vec![
2257 ColumnStatistics {
2258 distinct_count: Precision::Exact(5),
2259 max_value: Precision::Exact(ScalarValue::Int64(Some(21))),
2260 min_value: Precision::Exact(ScalarValue::Int64(Some(-4))),
2261 sum_value: Precision::Exact(ScalarValue::Int64(Some(42))),
2262 null_count: Precision::Exact(0),
2263 byte_size: Precision::Absent,
2264 },
2265 ColumnStatistics {
2266 distinct_count: Precision::Exact(1),
2267 max_value: Precision::Exact(ScalarValue::from("x")),
2268 min_value: Precision::Exact(ScalarValue::from("a")),
2269 sum_value: Precision::Absent,
2270 null_count: Precision::Exact(3),
2271 byte_size: Precision::Absent,
2272 },
2273 ColumnStatistics {
2274 distinct_count: Precision::Absent,
2275 max_value: Precision::Exact(ScalarValue::Float32(Some(1.1))),
2276 min_value: Precision::Exact(ScalarValue::Float32(Some(0.1))),
2277 sum_value: Precision::Exact(ScalarValue::Float32(Some(5.5))),
2278 null_count: Precision::Absent,
2279 byte_size: Precision::Absent,
2280 },
2281 ],
2282 }
2283 }
2284
2285 fn get_schema() -> Schema {
2286 let field_0 = Field::new("col0", DataType::Int64, false);
2287 let field_1 = Field::new("col1", DataType::Utf8, false);
2288 let field_2 = Field::new("col2", DataType::Float32, false);
2289 Schema::new(vec![field_0, field_1, field_2])
2290 }
2291
2292 #[test]
2293 fn test_projected_column_position_returns_output_position() {
2294 let projection = ProjectionExprs::new([
2295 ProjectionExpr::new(Arc::new(Column::new("col2", 2)), "col2"),
2296 ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "col0"),
2297 ]);
2298
2299 assert_eq!(
2300 projection.projected_column_position(&Column::new("col2", 2)),
2301 Some(0)
2302 );
2303 assert_eq!(
2304 projection.projected_column_position(&Column::new("col0", 0)),
2305 Some(1)
2306 );
2307 }
2308
2309 #[test]
2310 fn test_projected_column_position_returns_none_for_non_column_or_missing() {
2311 let projection = ProjectionExprs::new([
2312 ProjectionExpr::new(
2313 Arc::new(Literal::new(ScalarValue::Int64(Some(42)))),
2314 "col1",
2315 ),
2316 ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "col0"),
2317 ]);
2318
2319 assert_eq!(
2320 projection.projected_column_position(&Column::new("col1", 1)),
2321 None
2322 );
2323 assert_eq!(
2324 projection.projected_column_position(&Column::new("col2", 2)),
2325 None
2326 );
2327 }
2328
2329 #[test]
2330 fn test_stats_projection_columns_only() {
2331 let source = get_stats();
2332 let schema = get_schema();
2333
2334 let projection = ProjectionExprs::new(vec![
2335 ProjectionExpr {
2336 expr: Arc::new(Column::new("col1", 1)),
2337 alias: "col1".to_string(),
2338 },
2339 ProjectionExpr {
2340 expr: Arc::new(Column::new("col0", 0)),
2341 alias: "col0".to_string(),
2342 },
2343 ]);
2344
2345 let result = projection
2346 .project_statistics(source, &projection.project_schema(&schema).unwrap())
2347 .unwrap();
2348
2349 let expected = Statistics {
2350 num_rows: Precision::Exact(5),
2351 total_byte_size: Precision::Inexact(23),
2354 column_statistics: vec![
2355 ColumnStatistics {
2356 distinct_count: Precision::Exact(1),
2357 max_value: Precision::Exact(ScalarValue::from("x")),
2358 min_value: Precision::Exact(ScalarValue::from("a")),
2359 sum_value: Precision::Absent,
2360 null_count: Precision::Exact(3),
2361 byte_size: Precision::Absent,
2362 },
2363 ColumnStatistics {
2364 distinct_count: Precision::Exact(5),
2365 max_value: Precision::Exact(ScalarValue::Int64(Some(21))),
2366 min_value: Precision::Exact(ScalarValue::Int64(Some(-4))),
2367 sum_value: Precision::Exact(ScalarValue::Int64(Some(42))),
2368 null_count: Precision::Exact(0),
2369 byte_size: Precision::Absent,
2370 },
2371 ],
2372 };
2373
2374 assert_eq!(result, expected);
2375 }
2376
2377 #[test]
2378 fn test_stats_projection_column_with_primitive_width_only() {
2379 let source = get_stats();
2380 let schema = get_schema();
2381
2382 let projection = ProjectionExprs::new(vec![
2383 ProjectionExpr {
2384 expr: Arc::new(Column::new("col2", 2)),
2385 alias: "col2".to_string(),
2386 },
2387 ProjectionExpr {
2388 expr: Arc::new(Column::new("col0", 0)),
2389 alias: "col0".to_string(),
2390 },
2391 ]);
2392
2393 let result = projection
2394 .project_statistics(source, &projection.project_schema(&schema).unwrap())
2395 .unwrap();
2396
2397 let expected = Statistics {
2398 num_rows: Precision::Exact(5),
2399 total_byte_size: Precision::Exact(60),
2400 column_statistics: vec![
2401 ColumnStatistics {
2402 distinct_count: Precision::Absent,
2403 max_value: Precision::Exact(ScalarValue::Float32(Some(1.1))),
2404 min_value: Precision::Exact(ScalarValue::Float32(Some(0.1))),
2405 sum_value: Precision::Exact(ScalarValue::Float32(Some(5.5))),
2406 null_count: Precision::Absent,
2407 byte_size: Precision::Absent,
2408 },
2409 ColumnStatistics {
2410 distinct_count: Precision::Exact(5),
2411 max_value: Precision::Exact(ScalarValue::Int64(Some(21))),
2412 min_value: Precision::Exact(ScalarValue::Int64(Some(-4))),
2413 sum_value: Precision::Exact(ScalarValue::Int64(Some(42))),
2414 null_count: Precision::Exact(0),
2415 byte_size: Precision::Absent,
2416 },
2417 ],
2418 };
2419
2420 assert_eq!(result, expected);
2421 }
2422
2423 #[test]
2426 fn test_projection_new() -> Result<()> {
2427 let exprs = vec![
2428 ProjectionExpr {
2429 expr: Arc::new(Column::new("a", 0)),
2430 alias: "a".to_string(),
2431 },
2432 ProjectionExpr {
2433 expr: Arc::new(Column::new("b", 1)),
2434 alias: "b".to_string(),
2435 },
2436 ];
2437 let projection = ProjectionExprs::new(exprs.clone());
2438 assert_eq!(projection.as_ref().len(), 2);
2439 Ok(())
2440 }
2441
2442 #[test]
2443 fn test_projection_from_vec() -> Result<()> {
2444 let exprs = vec![ProjectionExpr {
2445 expr: Arc::new(Column::new("x", 0)),
2446 alias: "x".to_string(),
2447 }];
2448 let projection: ProjectionExprs = exprs.clone().into();
2449 assert_eq!(projection.as_ref().len(), 1);
2450 Ok(())
2451 }
2452
2453 #[test]
2454 fn test_projection_as_ref() -> Result<()> {
2455 let exprs = vec![
2456 ProjectionExpr {
2457 expr: Arc::new(Column::new("col1", 0)),
2458 alias: "col1".to_string(),
2459 },
2460 ProjectionExpr {
2461 expr: Arc::new(Column::new("col2", 1)),
2462 alias: "col2".to_string(),
2463 },
2464 ];
2465 let projection = ProjectionExprs::new(exprs);
2466 let as_ref: &[ProjectionExpr] = projection.as_ref();
2467 assert_eq!(as_ref.len(), 2);
2468 Ok(())
2469 }
2470
2471 #[test]
2472 fn test_column_indices_multiple_columns() -> Result<()> {
2473 let projection = ProjectionExprs::new(vec![
2475 ProjectionExpr {
2476 expr: Arc::new(Column::new("c", 5)),
2477 alias: "c".to_string(),
2478 },
2479 ProjectionExpr {
2480 expr: Arc::new(Column::new("b", 2)),
2481 alias: "b".to_string(),
2482 },
2483 ProjectionExpr {
2484 expr: Arc::new(Column::new("a", 0)),
2485 alias: "a".to_string(),
2486 },
2487 ]);
2488 assert_eq!(projection.column_indices(), vec![0, 2, 5]);
2490 Ok(())
2491 }
2492
2493 #[test]
2494 fn test_column_indices_duplicates() -> Result<()> {
2495 let projection = ProjectionExprs::new(vec![
2497 ProjectionExpr {
2498 expr: Arc::new(Column::new("a", 1)),
2499 alias: "a".to_string(),
2500 },
2501 ProjectionExpr {
2502 expr: Arc::new(Column::new("b", 3)),
2503 alias: "b".to_string(),
2504 },
2505 ProjectionExpr {
2506 expr: Arc::new(Column::new("a2", 1)), alias: "a2".to_string(),
2508 },
2509 ]);
2510 assert_eq!(projection.column_indices(), vec![1, 3]);
2511 Ok(())
2512 }
2513
2514 #[test]
2515 fn test_column_indices_unsorted() -> Result<()> {
2516 let projection = ProjectionExprs::new(vec![
2518 ProjectionExpr {
2519 expr: Arc::new(Column::new("c", 5)),
2520 alias: "c".to_string(),
2521 },
2522 ProjectionExpr {
2523 expr: Arc::new(Column::new("a", 1)),
2524 alias: "a".to_string(),
2525 },
2526 ProjectionExpr {
2527 expr: Arc::new(Column::new("b", 3)),
2528 alias: "b".to_string(),
2529 },
2530 ]);
2531 assert_eq!(projection.column_indices(), vec![1, 3, 5]);
2532 Ok(())
2533 }
2534
2535 #[test]
2536 fn test_column_indices_complex_expr() -> Result<()> {
2537 let expr = Arc::new(BinaryExpr::new(
2539 Arc::new(Column::new("a", 1)),
2540 Operator::Plus,
2541 Arc::new(Column::new("b", 4)),
2542 ));
2543 let projection = ProjectionExprs::new(vec![
2544 ProjectionExpr {
2545 expr,
2546 alias: "sum".to_string(),
2547 },
2548 ProjectionExpr {
2549 expr: Arc::new(Column::new("c", 2)),
2550 alias: "c".to_string(),
2551 },
2552 ]);
2553 assert_eq!(projection.column_indices(), vec![1, 2, 4]);
2555 Ok(())
2556 }
2557
2558 #[test]
2559 fn test_column_indices_empty() -> Result<()> {
2560 let projection = ProjectionExprs::new(vec![]);
2561 assert_eq!(projection.column_indices(), Vec::<usize>::new());
2562 Ok(())
2563 }
2564
2565 #[test]
2566 fn test_merge_simple_columns() -> Result<()> {
2567 let base_projection = ProjectionExprs::new(vec![
2569 ProjectionExpr {
2570 expr: Arc::new(Column::new("c", 2)),
2571 alias: "x".to_string(),
2572 },
2573 ProjectionExpr {
2574 expr: Arc::new(Column::new("b", 1)),
2575 alias: "y".to_string(),
2576 },
2577 ProjectionExpr {
2578 expr: Arc::new(Column::new("a", 0)),
2579 alias: "z".to_string(),
2580 },
2581 ]);
2582
2583 let top_projection = ProjectionExprs::new(vec![
2585 ProjectionExpr {
2586 expr: Arc::new(Column::new("y", 1)),
2587 alias: "col2".to_string(),
2588 },
2589 ProjectionExpr {
2590 expr: Arc::new(Column::new("x", 0)),
2591 alias: "col1".to_string(),
2592 },
2593 ]);
2594
2595 let merged = base_projection.try_merge(&top_projection)?;
2597 assert_snapshot!(format!("{merged}"), @"Projection[b@1 AS col2, c@2 AS col1]");
2598
2599 Ok(())
2600 }
2601
2602 #[test]
2603 fn test_merge_with_expressions() -> Result<()> {
2604 let base_projection = ProjectionExprs::new(vec![
2606 ProjectionExpr {
2607 expr: Arc::new(Column::new("c", 2)),
2608 alias: "x".to_string(),
2609 },
2610 ProjectionExpr {
2611 expr: Arc::new(Column::new("b", 1)),
2612 alias: "y".to_string(),
2613 },
2614 ProjectionExpr {
2615 expr: Arc::new(Column::new("a", 0)),
2616 alias: "z".to_string(),
2617 },
2618 ]);
2619
2620 let top_projection = ProjectionExprs::new(vec![
2622 ProjectionExpr {
2623 expr: Arc::new(BinaryExpr::new(
2624 Arc::new(Column::new("y", 1)),
2625 Operator::Plus,
2626 Arc::new(Column::new("z", 2)),
2627 )),
2628 alias: "c2".to_string(),
2629 },
2630 ProjectionExpr {
2631 expr: Arc::new(BinaryExpr::new(
2632 Arc::new(Column::new("x", 0)),
2633 Operator::Plus,
2634 Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
2635 )),
2636 alias: "c1".to_string(),
2637 },
2638 ]);
2639
2640 let merged = base_projection.try_merge(&top_projection)?;
2642 assert_snapshot!(format!("{merged}"), @"Projection[b@1 + a@0 AS c2, c@2 + 1 AS c1]");
2643
2644 Ok(())
2645 }
2646
2647 #[test]
2648 fn try_merge_error() {
2649 let base = ProjectionExprs::new(vec![
2651 ProjectionExpr {
2652 expr: Arc::new(Column::new("a", 0)),
2653 alias: "x".to_string(),
2654 },
2655 ProjectionExpr {
2656 expr: Arc::new(Column::new("b", 1)),
2657 alias: "y".to_string(),
2658 },
2659 ]);
2660
2661 let top = ProjectionExprs::new(vec![ProjectionExpr {
2663 expr: Arc::new(Column::new("z", 5)), alias: "result".to_string(),
2665 }]);
2666
2667 let err_msg = base.try_merge(&top).unwrap_err().to_string();
2669 assert!(
2670 err_msg.contains("Internal error: Column index 5 out of bounds for projected expressions of length 2"),
2671 "Unexpected error message: {err_msg}",
2672 );
2673 }
2674
2675 #[test]
2676 fn test_merge_empty_projection_with_literal() -> Result<()> {
2677 let base_projection = ProjectionExprs::new(vec![]);
2684
2685 let top_projection = ProjectionExprs::new(vec![ProjectionExpr {
2687 expr: Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
2688 alias: "Int64(1)".to_string(),
2689 }]);
2690
2691 let merged = base_projection.try_merge(&top_projection)?;
2694 assert_snapshot!(format!("{merged}"), @"Projection[1 AS Int64(1)]");
2695
2696 Ok(())
2697 }
2698
2699 #[test]
2700 fn test_update_expr_with_literal() -> Result<()> {
2701 let literal_expr: Arc<dyn PhysicalExpr> =
2703 Arc::new(Literal::new(ScalarValue::Int64(Some(42))));
2704 let empty_projection: Vec<ProjectionExpr> = vec![];
2705
2706 let result = update_expr(&literal_expr, &empty_projection, true)?;
2708 assert!(result.is_some(), "Literal expression should be valid");
2709
2710 let result_expr = result.unwrap();
2711 assert_eq!(
2712 result_expr.downcast_ref::<Literal>().unwrap().value(),
2713 &ScalarValue::Int64(Some(42))
2714 );
2715
2716 Ok(())
2717 }
2718
2719 #[test]
2720 fn test_update_expr_with_complex_literal_expr() -> Result<()> {
2721 let expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
2724 Arc::new(Literal::new(ScalarValue::Int64(Some(10)))),
2725 Operator::Plus,
2726 Arc::new(Column::new("x", 0)),
2727 ));
2728
2729 let base_projection = vec![ProjectionExpr {
2731 expr: Arc::new(Column::new("a", 5)),
2732 alias: "x".to_string(),
2733 }];
2734
2735 let result = update_expr(&expr, &base_projection, true)?;
2737 assert!(result.is_some(), "Expression should be valid");
2738
2739 let result_expr = result.unwrap();
2740 let binary = result_expr
2741 .downcast_ref::<BinaryExpr>()
2742 .expect("Should be a BinaryExpr");
2743
2744 assert!(binary.left().downcast_ref::<Literal>().is_some());
2746
2747 let right_col = binary
2749 .right()
2750 .downcast_ref::<Column>()
2751 .expect("Right should be a Column");
2752 assert_eq!(right_col.index(), 5);
2753
2754 Ok(())
2755 }
2756
2757 #[test]
2758 fn test_project_schema_simple_columns() -> Result<()> {
2759 let input_schema = get_schema();
2761
2762 let projection = ProjectionExprs::new(vec![
2764 ProjectionExpr {
2765 expr: Arc::new(Column::new("col2", 2)),
2766 alias: "c".to_string(),
2767 },
2768 ProjectionExpr {
2769 expr: Arc::new(Column::new("col0", 0)),
2770 alias: "a".to_string(),
2771 },
2772 ]);
2773
2774 let output_schema = projection.project_schema(&input_schema)?;
2775
2776 assert_eq!(output_schema.fields().len(), 2);
2778
2779 assert_eq!(output_schema.field(0).name(), "c");
2781 assert_eq!(output_schema.field(0).data_type(), &DataType::Float32);
2782
2783 assert_eq!(output_schema.field(1).name(), "a");
2785 assert_eq!(output_schema.field(1).data_type(), &DataType::Int64);
2786
2787 Ok(())
2788 }
2789
2790 #[test]
2791 fn test_project_schema_with_expressions() -> Result<()> {
2792 let input_schema = get_schema();
2794
2795 let projection = ProjectionExprs::new(vec![ProjectionExpr {
2797 expr: Arc::new(BinaryExpr::new(
2798 Arc::new(Column::new("col0", 0)),
2799 Operator::Plus,
2800 Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
2801 )),
2802 alias: "incremented".to_string(),
2803 }]);
2804
2805 let output_schema = projection.project_schema(&input_schema)?;
2806
2807 assert_eq!(output_schema.fields().len(), 1);
2809
2810 assert_eq!(output_schema.field(0).name(), "incremented");
2812 assert_eq!(output_schema.field(0).data_type(), &DataType::Int64);
2813
2814 Ok(())
2815 }
2816
2817 #[test]
2818 fn test_project_schema_preserves_metadata() -> Result<()> {
2819 let mut metadata = HashMap::new();
2821 metadata.insert("key".to_string(), "value".to_string());
2822 let field_with_metadata =
2823 Field::new("col0", DataType::Int64, false).with_metadata(metadata.clone());
2824 let input_schema = Schema::new(vec![
2825 field_with_metadata,
2826 Field::new("col1", DataType::Utf8, false),
2827 ]);
2828
2829 let projection = ProjectionExprs::new(vec![ProjectionExpr {
2831 expr: Arc::new(Column::new("col0", 0)),
2832 alias: "renamed".to_string(),
2833 }]);
2834
2835 let output_schema = projection.project_schema(&input_schema)?;
2836
2837 assert_eq!(output_schema.fields().len(), 1);
2839
2840 assert_eq!(output_schema.field(0).name(), "renamed");
2842 assert_eq!(output_schema.field(0).metadata(), &metadata);
2843
2844 Ok(())
2845 }
2846
2847 #[test]
2848 fn test_project_schema_empty() -> Result<()> {
2849 let input_schema = get_schema();
2850 let projection = ProjectionExprs::new(vec![]);
2851
2852 let output_schema = projection.project_schema(&input_schema)?;
2853
2854 assert_eq!(output_schema.fields().len(), 0);
2855
2856 Ok(())
2857 }
2858
2859 #[test]
2860 fn test_project_statistics_columns_only() -> Result<()> {
2861 let input_stats = get_stats();
2862 let input_schema = get_schema();
2863
2864 let projection = ProjectionExprs::new(vec![
2866 ProjectionExpr {
2867 expr: Arc::new(Column::new("col1", 1)),
2868 alias: "text".to_string(),
2869 },
2870 ProjectionExpr {
2871 expr: Arc::new(Column::new("col0", 0)),
2872 alias: "num".to_string(),
2873 },
2874 ]);
2875
2876 let output_stats = projection.project_statistics(
2877 input_stats,
2878 &projection.project_schema(&input_schema)?,
2879 )?;
2880
2881 assert_eq!(output_stats.num_rows, Precision::Exact(5));
2883
2884 assert_eq!(output_stats.column_statistics.len(), 2);
2886
2887 assert_eq!(
2889 output_stats.column_statistics[0].distinct_count,
2890 Precision::Exact(1)
2891 );
2892 assert_eq!(
2893 output_stats.column_statistics[0].max_value,
2894 Precision::Exact(ScalarValue::from("x"))
2895 );
2896
2897 assert_eq!(
2899 output_stats.column_statistics[1].distinct_count,
2900 Precision::Exact(5)
2901 );
2902 assert_eq!(
2903 output_stats.column_statistics[1].max_value,
2904 Precision::Exact(ScalarValue::Int64(Some(21)))
2905 );
2906
2907 Ok(())
2908 }
2909
2910 #[test]
2911 fn test_project_statistics_with_expressions() -> Result<()> {
2912 let input_stats = get_stats();
2913 let input_schema = get_schema();
2914
2915 let projection = ProjectionExprs::new(vec![
2917 ProjectionExpr {
2918 expr: Arc::new(BinaryExpr::new(
2919 Arc::new(Column::new("col0", 0)),
2920 Operator::Plus,
2921 Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
2922 )),
2923 alias: "incremented".to_string(),
2924 },
2925 ProjectionExpr {
2926 expr: Arc::new(Column::new("col1", 1)),
2927 alias: "text".to_string(),
2928 },
2929 ]);
2930
2931 let output_stats = projection.project_statistics(
2932 input_stats,
2933 &projection.project_schema(&input_schema)?,
2934 )?;
2935
2936 assert_eq!(output_stats.num_rows, Precision::Exact(5));
2938
2939 assert_eq!(output_stats.column_statistics.len(), 2);
2941
2942 assert_eq!(
2944 output_stats.column_statistics[0].distinct_count,
2945 Precision::Absent
2946 );
2947 assert_eq!(
2948 output_stats.column_statistics[0].max_value,
2949 Precision::Absent
2950 );
2951
2952 assert_eq!(
2954 output_stats.column_statistics[1].distinct_count,
2955 Precision::Exact(1)
2956 );
2957
2958 Ok(())
2959 }
2960
2961 #[test]
2962 fn test_project_statistics_with_same_type_cast_is_exact_passthrough() -> Result<()> {
2963 let input_stats = get_stats();
2968 let col0_stats = input_stats.column_statistics[0].clone();
2969 let input_schema = get_schema();
2970
2971 let projection = ProjectionExprs::new(vec![ProjectionExpr {
2972 expr: Arc::new(CastExpr::new(
2973 Arc::new(Column::new("col0", 0)),
2974 DataType::Int64,
2975 None,
2976 )),
2977 alias: "casted".to_string(),
2978 }]);
2979
2980 let output_stats = projection.project_statistics(
2981 input_stats,
2982 &projection.project_schema(&input_schema)?,
2983 )?;
2984
2985 assert_eq!(output_stats.column_statistics[0], col0_stats);
2986
2987 Ok(())
2988 }
2989
2990 #[test]
2991 fn test_project_statistics_with_cast() -> Result<()> {
2992 let input_stats = get_stats();
2993 let input_schema = get_schema();
2994
2995 let projection = ProjectionExprs::new(vec![ProjectionExpr {
2997 expr: Arc::new(CastExpr::new(
2998 Arc::new(Column::new("col0", 0)),
2999 DataType::Int32,
3000 None,
3001 )),
3002 alias: "casted".to_string(),
3003 }]);
3004
3005 let output_stats = projection.project_statistics(
3006 input_stats,
3007 &projection.project_schema(&input_schema)?,
3008 )?;
3009
3010 assert_eq!(
3011 output_stats.column_statistics[0].min_value,
3012 Precision::Exact(ScalarValue::Int32(Some(-4)))
3013 );
3014 assert_eq!(
3015 output_stats.column_statistics[0].max_value,
3016 Precision::Exact(ScalarValue::Int32(Some(21)))
3017 );
3018
3019 Ok(())
3020 }
3021
3022 #[test]
3023 fn test_project_statistics_duplicate_column() -> Result<()> {
3024 let input_stats = get_stats();
3025 let col0 = input_stats.column_statistics[0].clone();
3026 let projection = ProjectionExprs::new([
3027 ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "a"),
3028 ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "b"),
3029 ]);
3030
3031 let output_schema = projection.project_schema(&get_schema())?;
3032 let output_stats = projection.project_statistics(input_stats, &output_schema)?;
3033
3034 assert_eq!(output_stats.column_statistics, vec![col0.clone(), col0]);
3035 Ok(())
3036 }
3037
3038 #[test]
3039 fn test_project_statistics_column_and_cast() -> Result<()> {
3040 let input_stats = get_stats();
3041 let col0 = input_stats.column_statistics[0].clone();
3042 let projection = ProjectionExprs::new([
3043 ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "num"),
3044 ProjectionExpr::new(
3045 Arc::new(CastExpr::new(
3046 Arc::new(Column::new("col0", 0)),
3047 DataType::Int32,
3048 None,
3049 )),
3050 "casted",
3051 ),
3052 ]);
3053
3054 let output_schema = projection.project_schema(&get_schema())?;
3055 let output_stats = projection.project_statistics(input_stats, &output_schema)?;
3056
3057 assert_eq!(output_stats.column_statistics[0], col0);
3058 assert_eq!(
3059 output_stats.column_statistics[1],
3060 ColumnStatistics {
3061 min_value: Precision::Exact(ScalarValue::Int32(Some(-4))),
3062 max_value: Precision::Exact(ScalarValue::Int32(Some(21))),
3063 distinct_count: Precision::Exact(5),
3064 null_count: Precision::Exact(0),
3065 sum_value: Precision::Absent,
3066 byte_size: Precision::Absent,
3067 }
3068 );
3069
3070 Ok(())
3071 }
3072
3073 #[test]
3074 fn test_project_statistics_missing_column_stats_are_unknown() -> Result<()> {
3075 let mut input_stats = get_stats();
3076 let input_schema = get_schema();
3077 input_stats.column_statistics.truncate(2);
3078
3079 let projection = ProjectionExprs::new(vec![
3083 ProjectionExpr {
3084 expr: Arc::new(Column::new("col2", 2)),
3085 alias: "virtual_col".to_string(),
3086 },
3087 ProjectionExpr {
3088 expr: Arc::new(CastExpr::new(
3089 Arc::new(Column::new("col2", 2)),
3090 DataType::Float64,
3091 None,
3092 )),
3093 alias: "casted_virtual_col".to_string(),
3094 },
3095 ProjectionExpr {
3096 expr: Arc::new(Column::new("col0", 0)),
3097 alias: "physical_col".to_string(),
3098 },
3099 ]);
3100
3101 let output_stats = projection.project_statistics(
3102 input_stats,
3103 &projection.project_schema(&input_schema)?,
3104 )?;
3105
3106 assert_eq!(output_stats.column_statistics.len(), 3);
3107 assert_eq!(
3108 output_stats.column_statistics[0],
3109 ColumnStatistics::new_unknown()
3110 );
3111 assert_eq!(
3112 output_stats.column_statistics[1],
3113 ColumnStatistics::new_unknown()
3114 );
3115 assert_eq!(
3116 output_stats.column_statistics[2].max_value,
3117 Precision::Exact(ScalarValue::Int64(Some(21)))
3118 );
3119
3120 Ok(())
3121 }
3122
3123 #[test]
3124 fn test_project_statistics_primitive_width_only() -> Result<()> {
3125 let input_stats = get_stats();
3126 let input_schema = get_schema();
3127
3128 let projection = ProjectionExprs::new(vec![
3130 ProjectionExpr {
3131 expr: Arc::new(Column::new("col2", 2)),
3132 alias: "f".to_string(),
3133 },
3134 ProjectionExpr {
3135 expr: Arc::new(Column::new("col0", 0)),
3136 alias: "i".to_string(),
3137 },
3138 ]);
3139
3140 let output_stats = projection.project_statistics(
3141 input_stats,
3142 &projection.project_schema(&input_schema)?,
3143 )?;
3144
3145 assert_eq!(output_stats.num_rows, Precision::Exact(5));
3147
3148 assert_eq!(output_stats.total_byte_size, Precision::Exact(60));
3151
3152 assert_eq!(output_stats.column_statistics.len(), 2);
3154
3155 Ok(())
3156 }
3157
3158 #[test]
3159 fn test_project_statistics_empty() -> Result<()> {
3160 let input_stats = get_stats();
3161 let input_schema = get_schema();
3162
3163 let projection = ProjectionExprs::new(vec![]);
3164
3165 let output_stats = projection.project_statistics(
3166 input_stats,
3167 &projection.project_schema(&input_schema)?,
3168 )?;
3169
3170 assert_eq!(output_stats.num_rows, Precision::Exact(5));
3172
3173 assert_eq!(output_stats.column_statistics.len(), 0);
3175
3176 assert_eq!(output_stats.total_byte_size, Precision::Exact(0));
3178
3179 Ok(())
3180 }
3181
3182 #[test]
3184 fn test_project_statistics_with_literal() -> Result<()> {
3185 let input_stats = get_stats();
3186 let input_schema = get_schema();
3187
3188 let projection = ProjectionExprs::new(vec![
3190 ProjectionExpr {
3191 expr: Arc::new(Literal::new(ScalarValue::Int64(Some(42)))),
3192 alias: "constant".to_string(),
3193 },
3194 ProjectionExpr {
3195 expr: Arc::new(Column::new("col0", 0)),
3196 alias: "num".to_string(),
3197 },
3198 ]);
3199
3200 let output_stats = projection.project_statistics(
3201 input_stats,
3202 &projection.project_schema(&input_schema)?,
3203 )?;
3204
3205 assert_eq!(output_stats.num_rows, Precision::Exact(5));
3207
3208 assert_eq!(output_stats.column_statistics.len(), 2);
3210
3211 assert_eq!(
3213 output_stats.column_statistics[0].min_value,
3214 Precision::Exact(ScalarValue::Int64(Some(42)))
3215 );
3216 assert_eq!(
3217 output_stats.column_statistics[0].max_value,
3218 Precision::Exact(ScalarValue::Int64(Some(42)))
3219 );
3220 assert_eq!(
3221 output_stats.column_statistics[0].distinct_count,
3222 Precision::Exact(1)
3223 );
3224 assert_eq!(
3225 output_stats.column_statistics[0].null_count,
3226 Precision::Exact(0)
3227 );
3228 assert_eq!(
3230 output_stats.column_statistics[0].byte_size,
3231 Precision::Exact(40)
3232 );
3233 assert_eq!(
3235 output_stats.column_statistics[0].sum_value,
3236 Precision::Exact(ScalarValue::Int64(Some(210)))
3237 );
3238
3239 assert_eq!(
3241 output_stats.column_statistics[1].distinct_count,
3242 Precision::Exact(5)
3243 );
3244 assert_eq!(
3245 output_stats.column_statistics[1].max_value,
3246 Precision::Exact(ScalarValue::Int64(Some(21)))
3247 );
3248
3249 Ok(())
3250 }
3251
3252 #[test]
3253 fn test_project_statistics_with_i32_literal_sum_widens_to_i64() -> Result<()> {
3254 let input_stats = get_stats();
3255 let input_schema = get_schema();
3256
3257 let projection = ProjectionExprs::new(vec![
3258 ProjectionExpr {
3259 expr: Arc::new(Literal::new(ScalarValue::Int32(Some(10)))),
3260 alias: "constant".to_string(),
3261 },
3262 ProjectionExpr {
3263 expr: Arc::new(Column::new("col0", 0)),
3264 alias: "num".to_string(),
3265 },
3266 ]);
3267
3268 let output_stats = projection.project_statistics(
3269 input_stats,
3270 &projection.project_schema(&input_schema)?,
3271 )?;
3272
3273 assert_eq!(
3274 output_stats.column_statistics[0].sum_value,
3275 Precision::Exact(ScalarValue::Int64(Some(50)))
3276 );
3277
3278 Ok(())
3279 }
3280
3281 #[test]
3283 fn test_project_statistics_with_null_literal() -> Result<()> {
3284 let input_stats = get_stats();
3285 let input_schema = get_schema();
3286
3287 let projection = ProjectionExprs::new(vec![
3289 ProjectionExpr {
3290 expr: Arc::new(Literal::new(ScalarValue::Int64(None))),
3291 alias: "null_col".to_string(),
3292 },
3293 ProjectionExpr {
3294 expr: Arc::new(Column::new("col0", 0)),
3295 alias: "num".to_string(),
3296 },
3297 ]);
3298
3299 let output_stats = projection.project_statistics(
3300 input_stats,
3301 &projection.project_schema(&input_schema)?,
3302 )?;
3303
3304 assert_eq!(output_stats.num_rows, Precision::Exact(5));
3306
3307 assert_eq!(output_stats.column_statistics.len(), 2);
3309
3310 assert_eq!(
3312 output_stats.column_statistics[0].min_value,
3313 Precision::Exact(ScalarValue::Int64(None))
3314 );
3315 assert_eq!(
3316 output_stats.column_statistics[0].max_value,
3317 Precision::Exact(ScalarValue::Int64(None))
3318 );
3319 assert_eq!(
3320 output_stats.column_statistics[0].distinct_count,
3321 Precision::Exact(1) );
3323 assert_eq!(
3324 output_stats.column_statistics[0].null_count,
3325 Precision::Exact(5) );
3327 assert_eq!(
3328 output_stats.column_statistics[0].byte_size,
3329 Precision::Exact(0)
3330 );
3331 assert_eq!(
3332 output_stats.column_statistics[0].sum_value,
3333 Precision::Exact(ScalarValue::Int64(None))
3334 );
3335
3336 assert_eq!(
3338 output_stats.column_statistics[1].distinct_count,
3339 Precision::Exact(5)
3340 );
3341 assert_eq!(
3342 output_stats.column_statistics[1].max_value,
3343 Precision::Exact(ScalarValue::Int64(Some(21)))
3344 );
3345
3346 Ok(())
3347 }
3348
3349 #[test]
3351 fn test_project_statistics_with_complex_type_literal() -> Result<()> {
3352 let input_stats = get_stats();
3353 let input_schema = get_schema();
3354
3355 let projection = ProjectionExprs::new(vec![
3357 ProjectionExpr {
3358 expr: Arc::new(Literal::new(ScalarValue::Utf8(Some(
3359 "hello".to_string(),
3360 )))),
3361 alias: "text".to_string(),
3362 },
3363 ProjectionExpr {
3364 expr: Arc::new(Column::new("col0", 0)),
3365 alias: "num".to_string(),
3366 },
3367 ]);
3368
3369 let output_stats = projection.project_statistics(
3370 input_stats,
3371 &projection.project_schema(&input_schema)?,
3372 )?;
3373
3374 assert_eq!(output_stats.num_rows, Precision::Exact(5));
3376
3377 assert_eq!(output_stats.column_statistics.len(), 2);
3379
3380 assert_eq!(
3383 output_stats.column_statistics[0].min_value,
3384 Precision::Exact(ScalarValue::Utf8(Some("hello".to_string())))
3385 );
3386 assert_eq!(
3387 output_stats.column_statistics[0].max_value,
3388 Precision::Exact(ScalarValue::Utf8(Some("hello".to_string())))
3389 );
3390 assert_eq!(
3391 output_stats.column_statistics[0].distinct_count,
3392 Precision::Exact(1)
3393 );
3394 assert_eq!(
3395 output_stats.column_statistics[0].null_count,
3396 Precision::Exact(0)
3397 );
3398 assert_eq!(
3401 output_stats.column_statistics[0].byte_size,
3402 Precision::Absent
3403 );
3404 assert_eq!(
3407 output_stats.column_statistics[0].sum_value,
3408 Precision::Absent
3409 );
3410
3411 assert_eq!(
3413 output_stats.column_statistics[1].distinct_count,
3414 Precision::Exact(5)
3415 );
3416 assert_eq!(
3417 output_stats.column_statistics[1].max_value,
3418 Precision::Exact(ScalarValue::Int64(Some(21)))
3419 );
3420
3421 Ok(())
3422 }
3423}