@@ -725,6 +725,60 @@ impl ProjectionExprs {
725725 stats. column_statistics = column_statistics;
726726 Ok ( stats)
727727 }
728+
729+ /// Returns the output position of `column` if this projection contains it.
730+ ///
731+ /// This only matches projection expressions that are exactly [`Column`] expressions.
732+ /// Computed expressions, even if they reference `column`, do not match. The
733+ /// comparison uses [`Column`] equality, so both the name and index must match.
734+ /// If the same column appears more than once, this returns the first matching
735+ /// position.
736+ ///
737+ /// # Example
738+ ///
739+ /// ```rust
740+ /// use datafusion_common::ScalarValue;
741+ /// use datafusion_physical_expr::expressions::{Column, Literal};
742+ /// use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs};
743+ /// use std::sync::Arc;
744+ ///
745+ /// let projection = ProjectionExprs::new([
746+ /// ProjectionExpr::new(Arc::new(Column::new("b", 1)), "b"),
747+ /// ProjectionExpr::new(
748+ /// Arc::new(Literal::new(ScalarValue::Int32(Some(42)))),
749+ /// "answer",
750+ /// ),
751+ /// ProjectionExpr::new(Arc::new(Column::new("a", 0)), "a"),
752+ /// ]);
753+ ///
754+ /// assert_eq!(
755+ /// projection.projected_column_position(&Column::new("b", 1)),
756+ /// Some(0)
757+ /// );
758+ /// assert_eq!(
759+ /// projection.projected_column_position(&Column::new("a", 0)),
760+ /// Some(2)
761+ /// );
762+ ///
763+ /// // The literal projection is not a Column expression.
764+ /// assert_eq!(
765+ /// projection.projected_column_position(&Column::new("answer", 1)),
766+ /// None
767+ /// );
768+ ///
769+ /// // Columns not present in the projection also return None.
770+ /// assert_eq!(
771+ /// projection.projected_column_position(&Column::new("c", 2)),
772+ /// None
773+ /// );
774+ /// ```
775+ pub fn projected_column_position ( & self , column : & Column ) -> Option < usize > {
776+ self . iter ( ) . position ( |expr| {
777+ expr. expr
778+ . downcast_ref :: < Column > ( )
779+ . is_some_and ( |projected| projected == column)
780+ } )
781+ }
728782}
729783
730784/// Propagate column statistics through CAST projections. Other expressions
@@ -2212,6 +2266,43 @@ pub(crate) mod tests {
22122266 Schema :: new ( vec ! [ field_0, field_1, field_2] )
22132267 }
22142268
2269+ #[ test]
2270+ fn test_projected_column_position_returns_output_position ( ) {
2271+ let projection = ProjectionExprs :: new ( [
2272+ ProjectionExpr :: new ( Arc :: new ( Column :: new ( "col2" , 2 ) ) , "col2" ) ,
2273+ ProjectionExpr :: new ( Arc :: new ( Column :: new ( "col0" , 0 ) ) , "col0" ) ,
2274+ ] ) ;
2275+
2276+ assert_eq ! (
2277+ projection. projected_column_position( & Column :: new( "col2" , 2 ) ) ,
2278+ Some ( 0 )
2279+ ) ;
2280+ assert_eq ! (
2281+ projection. projected_column_position( & Column :: new( "col0" , 0 ) ) ,
2282+ Some ( 1 )
2283+ ) ;
2284+ }
2285+
2286+ #[ test]
2287+ fn test_projected_column_position_returns_none_for_non_column_or_missing ( ) {
2288+ let projection = ProjectionExprs :: new ( [
2289+ ProjectionExpr :: new (
2290+ Arc :: new ( Literal :: new ( ScalarValue :: Int64 ( Some ( 42 ) ) ) ) ,
2291+ "col1" ,
2292+ ) ,
2293+ ProjectionExpr :: new ( Arc :: new ( Column :: new ( "col0" , 0 ) ) , "col0" ) ,
2294+ ] ) ;
2295+
2296+ assert_eq ! (
2297+ projection. projected_column_position( & Column :: new( "col1" , 1 ) ) ,
2298+ None
2299+ ) ;
2300+ assert_eq ! (
2301+ projection. projected_column_position( & Column :: new( "col2" , 2 ) ) ,
2302+ None
2303+ ) ;
2304+ }
2305+
22152306 #[ test]
22162307 fn test_stats_projection_columns_only ( ) {
22172308 let source = get_stats ( ) ;
0 commit comments