feat: Add new input_file_name UDF for file-backed scans

datafusion/datasource-parquet/src/opener/mod.rs

@@ -38,6 +38,7 @@ use arrow::datatypes::DataType;
 use datafusion_datasource::morsel::{Morsel, MorselPlan, MorselPlanner, Morselizer};
 use datafusion_physical_expr::projection::ProjectionExprs;
 use datafusion_physical_expr_adapter::replace_columns_with_literals;
+use datafusion_physical_expr_adapter::schema_rewriter::rewrite_input_file_name_in_projection;
 use std::collections::{HashMap, VecDeque};
 use std::fmt;
 use std::future::Future;
@@ -770,6 +771,9 @@ impl ParquetMorselizer {
                 .transpose()?;
         }
 
+        // Replace any `input_file_name()` UDFs in the projection with a literal for this file.
+        projection = rewrite_input_file_name_in_projection(projection, &file_name)?;
+
         let predicate_creation_errors = MetricBuilder::new(&self.metrics)
             .with_category(MetricCategory::Rows)
             .global_counter("num_predicate_creation_errors");
@@ -2970,8 +2974,12 @@ mod test {
     /// (e.g. `row_number`) plumbed through `TableSchema`/`ParquetOpener`.
     mod virtual_columns {
         use super::*;
-        use arrow::array::{Array, Int64Array};
+        use arrow::array::{Array, Int64Array, StringArray};
         use arrow::datatypes::FieldRef;
+        use datafusion_common::config::ConfigOptions;
+        use datafusion_expr::ScalarUDF;
+        use datafusion_functions::core::input_file_name::InputFileNameFunc;
+        use datafusion_physical_expr::{ScalarFunctionExpr, projection::ProjectionExpr};
         use parquet::arrow::RowNumber;
 
         /// Build a parquet `row_number` virtual column field. Spark's
@@ -2985,6 +2993,16 @@ mod test {
             )
         }
 
+        fn input_file_name_expr() -> Arc<dyn PhysicalExpr> {
+            Arc::new(ScalarFunctionExpr::new(
+                "input_file_name",
+                Arc::new(ScalarUDF::from(InputFileNameFunc::new())),
+                vec![],
+                Arc::new(Field::new("input_file_name", DataType::Utf8, true)),
+                Arc::new(ConfigOptions::default()),
+            ))
+        }
+
         /// Collect every `Int64` value from the given column in every batch
         /// of a stream. Used to verify the `row_number` column end to end.
         async fn collect_int64_values(
@@ -3104,6 +3122,44 @@ mod test {
             assert_eq!(row_numbers, vec![0, 1, 2, 3]);
         }
 
+        #[tokio::test]
+        async fn test_input_file_name_projection() {
+            let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+            let path = "dir/input_file_name.parquet";
+            let (file_schema, data_size) = write_grouped_file(&store, path, 1, 3).await;
+
+            let projection = ProjectionExprs::new([
+                ProjectionExpr::new(Arc::new(Column::new("value", 0)), "value"),
+                ProjectionExpr::new(input_file_name_expr(), "file_name"),
+            ]);
+
+            let morselizer = ParquetMorselizerBuilder::new()
+                .with_store(Arc::clone(&store))
+                .with_schema(file_schema)
+                .with_projection(projection)
+                .build();
+
+            let file =
+                PartitionedFile::new(path.to_string(), u64::try_from(data_size).unwrap());
+            let mut stream = open_file(&morselizer, file).await.unwrap();
+            let batch = stream.next().await.unwrap().unwrap();
+            assert!(stream.next().await.is_none());
+
+            assert_eq!(batch.num_columns(), 2);
+            assert_eq!(batch.schema().field(0).name(), "value");
+            assert_eq!(batch.schema().field(1).name(), "file_name");
+
+            let file_names = batch
+                .column(1)
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .expect("file_name column should be Utf8");
+            assert_eq!(file_names.len(), 3);
+            for i in 0..file_names.len() {
+                assert_eq!(file_names.value(i), path);
+            }
+        }
+
         #[tokio::test]
         async fn test_row_index_multi_row_group() {
             let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;

datafusion/datasource/Cargo.toml

@@ -74,6 +74,7 @@ zstd = { workspace = true, optional = true }
 
 [dev-dependencies]
 criterion = { workspace = true }
+datafusion-functions = { workspace = true }
 insta = { workspace = true }
 tempfile = { workspace = true }
 

datafusion/datasource/src/projection.rs

@@ -26,6 +26,7 @@ use datafusion_physical_expr::{
     expressions::{Column, Literal},
     projection::{ProjectionExpr, ProjectionExprs},
 };
+use datafusion_physical_expr_adapter::schema_rewriter::rewrite_input_file_name_in_projection;
 use futures::{FutureExt, StreamExt};
 use itertools::Itertools;
 
@@ -69,6 +70,7 @@ impl ProjectionOpener {
 impl FileOpener for ProjectionOpener {
     fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
         let partition_values = partitioned_file.partition_values.clone();
+
         // Modify any references to partition columns in the projection expressions
         // and substitute them with literal values from PartitionedFile.partition_values
         let projection = if self.partition_columns.is_empty() {
@@ -80,6 +82,11 @@ impl FileOpener for ProjectionOpener {
                 partition_values,
             )
         };
+        // Replace `input_file_name()` with a per-file literal if present.
+        let projection = rewrite_input_file_name_in_projection(
+            projection,
+            partitioned_file.object_meta.location.as_ref(),
+        )?;
         let projector = projection.make_projector(&self.input_schema)?;
 
         let inner = self.inner.open(partitioned_file)?;
@@ -287,15 +294,31 @@ impl SplitProjection {
 mod test {
     use std::sync::Arc;
 
-    use arrow::array::AsArray;
-    use arrow::datatypes::{DataType, SchemaRef};
-    use datafusion_common::{DFSchema, ScalarValue, record_batch};
-    use datafusion_expr::{Expr, col, execution_props::ExecutionProps};
-    use datafusion_physical_expr::{create_physical_exprs, projection::ProjectionExpr};
+    use arrow::array::{AsArray, RecordBatch};
+    use arrow::datatypes::{DataType, Field, SchemaRef};
+    use datafusion_common::{DFSchema, ScalarValue, config::ConfigOptions, record_batch};
+    use datafusion_expr::{Expr, ScalarUDF, col, execution_props::ExecutionProps};
+    use datafusion_functions::core::input_file_name::InputFileNameFunc;
+    use datafusion_physical_expr::{
+        ScalarFunctionExpr, create_physical_exprs, projection::ProjectionExpr,
+    };
+    use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+    use futures::{FutureExt, StreamExt};
     use itertools::Itertools;
 
     use super::*;
 
+    struct StaticBatchOpener {
+        batch: RecordBatch,
+    }
+
+    impl FileOpener for StaticBatchOpener {
+        fn open(&self, _partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
+            let batch = self.batch.clone();
+            Ok(async move { Ok(futures::stream::iter([Ok(batch)]).boxed()) }.boxed())
+        }
+    }
+
     fn create_projection_exprs<'a>(
         exprs: impl IntoIterator<Item = &'a Expr>,
         schema: &SchemaRef,
@@ -311,6 +334,68 @@ mod test {
         ProjectionExprs::from(projection_exprs)
     }
 
+    fn input_file_name_expr() -> Arc<dyn PhysicalExpr> {
+        Arc::new(ScalarFunctionExpr::new(
+            "input_file_name",
+            Arc::new(ScalarUDF::from(InputFileNameFunc::new())),
+            vec![],
+            Arc::new(Field::new("input_file_name", DataType::Utf8, true)),
+            Arc::new(ConfigOptions::default()),
+        ))
+    }
+
+    #[tokio::test]
+    async fn test_projection_opener_rewrites_input_file_name_with_partitions() {
+        let file_schema = Schema::new(vec![Field::new("value", DataType::Int32, false)]);
+        let projection = ProjectionExprs::new([
+            ProjectionExpr::new(Arc::new(Column::new("value", 0)), "value"),
+            ProjectionExpr::new(Arc::new(Column::new("part", 1)), "part"),
+            ProjectionExpr::new(input_file_name_expr(), "file_name"),
+        ]);
+        let split = SplitProjection::new(&file_schema, &projection);
+        let input_batch =
+            record_batch!(("value", Int32, vec![10, 20])).expect("input batch");
+
+        let opener = ProjectionOpener::try_new(
+            split,
+            Arc::new(StaticBatchOpener { batch: input_batch }),
+            &file_schema,
+        )
+        .expect("projection opener");
+
+        let mut file = PartitionedFile::new("part=west/data.csv", 100);
+        file.partition_values = vec![ScalarValue::from("west")];
+        let mut stream = opener
+            .open(file)
+            .expect("open projection")
+            .await
+            .expect("inner stream");
+        let batch = stream
+            .next()
+            .await
+            .expect("one projected batch")
+            .expect("projected batch");
+        assert!(stream.next().await.is_none());
+
+        assert_eq!(batch.schema().field(0).name(), "value");
+        assert_eq!(batch.schema().field(1).name(), "part");
+        assert_eq!(batch.schema().field(2).name(), "file_name");
+
+        let values = batch
+            .column(0)
+            .as_primitive::<arrow::datatypes::Int32Type>();
+        assert_eq!(values.value(0), 10);
+        assert_eq!(values.value(1), 20);
+
+        let parts = batch.column(1).as_string::<i32>();
+        assert_eq!(parts.value(0), "west");
+        assert_eq!(parts.value(1), "west");
+
+        let file_names = batch.column(2).as_string::<i32>();
+        assert_eq!(file_names.value(0), "part=west/data.csv");
+        assert_eq!(file_names.value(1), "part=west/data.csv");
+    }
+
     #[test]
     fn test_split_projection_with_partition_columns() {
         use arrow::array::AsArray;

datafusion/functions/src/core/input_file_name.rs

@@ -0,0 +1,95 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`InputFileNameFunc`]: Implementation of the `input_file_name` function.
+
+use arrow::datatypes::DataType;
+use datafusion_common::{exec_err, utils::take_function_args};
+use datafusion_doc::Documentation;
+use datafusion_expr::{
+    ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+    Volatility,
+};
+use datafusion_macros::user_doc;
+
+#[user_doc(
+    doc_section(label = "Other Functions"),
+    description = r#"Returns the path of the input file that produced the current row.
+
+Note: file paths/URIs may be sensitive metadata depending on your environment.
+
+This function is intended to be rewritten at file-scan time (when the file is
+known). If the input file is not known (for example, if this function is
+evaluated outside a file scan, or was not pushed down into one), direct evaluation returns an error.
+"#,
+    syntax_example = "input_file_name()",
+    sql_example = r#"```sql
+SELECT input_file_name() FROM t;
+```"#
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct InputFileNameFunc {
+    signature: Signature,
+}
+
+impl Default for InputFileNameFunc {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl InputFileNameFunc {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::nullary(Volatility::Volatile),
+        }
+    }
+}
+
+impl ScalarUDFImpl for InputFileNameFunc {
+    fn name(&self) -> &str {
+        "input_file_name"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
+        let [] = take_function_args(self.name(), arg_types)?;
+        Ok(DataType::Utf8)
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: ScalarFunctionArgs,
+    ) -> datafusion_common::Result<ColumnarValue> {
+        let [] = take_function_args(self.name(), args.args)?;
+
+        exec_err!(
+            "input_file_name() is source dependent and cannot be evaluated directly"
+        )
+    }
+
+    fn placement(&self, _args: &[ExpressionPlacement]) -> ExpressionPlacement {
+        ExpressionPlacement::MoveTowardsLeafNodes
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.doc()
+    }
+}

datafusion/functions/src/core/mod.rs

@@ -32,6 +32,7 @@ pub mod file_row_index;
 pub mod getfield;
 pub mod greatest;
 mod greatest_least_utils;
+pub mod input_file_name;
 pub mod least;
 pub mod named_struct;
 pub mod nullif;
@@ -69,6 +70,7 @@ make_udf_function!(arrow_metadata::ArrowMetadataFunc, arrow_metadata);
 make_udf_function!(with_metadata::WithMetadataFunc, with_metadata);
 make_udf_function!(arrow_field::ArrowFieldFunc, arrow_field);
 make_udf_function!(file_row_index::FileRowIndexFunc, file_row_index);
+make_udf_function!(input_file_name::InputFileNameFunc, input_file_name);
 
 pub mod expr_fn {
     use datafusion_expr::{Expr, Literal};
@@ -117,7 +119,12 @@ pub mod expr_fn {
         arrow_metadata,
         "Returns the metadata of the input expression",
         args,
-    ),(
+    ),
+    (
+        input_file_name,
+        "Returns the path of the input file that produced the current row",
+    ),
+    (
         with_metadata,
         "Attaches Arrow field metadata (key/value pairs) to the input expression",
         args,
@@ -200,6 +207,7 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
         union_extract(),
         union_tag(),
         version(),
+        input_file_name(),
         r#struct(),
         file_row_index(),
     ]