Add file_row_index UDF to query file-level row indexes from Parquet files

datafusion/datasource-parquet/src/row_filter.rs

@@ -72,6 +72,7 @@ use arrow::array::BooleanArray;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::error::{ArrowError, Result as ArrowResult};
 use arrow::record_batch::RecordBatch;
+use datafusion_functions::core::file_row_index::FileRowIndexFunc;
 use datafusion_functions::core::getfield::GetFieldFunc;
 use parquet::arrow::ProjectionMask;
 use parquet::arrow::arrow_reader::{ArrowPredicate, RowFilter};
@@ -260,6 +261,9 @@ struct PushdownChecker<'schema> {
     non_primitive_columns: bool,
     /// Does the expression reference any columns not present in the file schema?
     projected_columns: bool,
+    /// Does the expression references a ScalarUDF that requires some rewrite
+    /// and therefore can't be pushed down into the row-filter.
+    has_unpushable_udfs: bool,
     /// Indices into the file schema of columns required to evaluate the expression.
     /// Does not include struct columns accessed via `get_field`.
     required_columns: Vec<usize>,
@@ -276,6 +280,7 @@ impl<'schema> PushdownChecker<'schema> {
         Self {
             non_primitive_columns: false,
             projected_columns: false,
+            has_unpushable_udfs: false,
             required_columns: Vec::new(),
             struct_field_accesses: Vec::new(),
             allow_list_columns,
@@ -372,7 +377,7 @@ impl<'schema> PushdownChecker<'schema> {
 
     #[inline]
     fn prevents_pushdown(&self) -> bool {
-        self.non_primitive_columns || self.projected_columns
+        self.non_primitive_columns || self.projected_columns || self.has_unpushable_udfs
     }
 
     /// Consumes the checker and returns sorted, deduplicated column indices
@@ -484,6 +489,13 @@ impl TreeNodeVisitor<'_> for PushdownChecker<'_> {
             return Ok(recursion);
         }
 
+        if ScalarFunctionExpr::try_downcast_func::<FileRowIndexFunc>(node.as_ref())
+            .is_some()
+        {
+            self.has_unpushable_udfs = true;
+            return Ok(TreeNodeRecursion::Jump);
+        }
+
         Ok(TreeNodeRecursion::Continue)
     }
 }

datafusion/datasource-parquet/src/source.rs

@@ -26,6 +26,9 @@ use crate::opener::ParquetMorselizer;
 use crate::opener::build_pruning_predicates;
 use crate::opener::build_virtual_columns_state;
 use crate::row_filter::can_expr_be_pushed_down_with_schemas;
+use arrow_schema::Fields;
+use arrow_schema::extension::ExtensionType;
+use arrow_schema::{DataType, Field};
 use datafusion_common::config::ConfigOptions;
 #[cfg(feature = "parquet_encryption")]
 use datafusion_common::config::EncryptionFactoryOptions;
@@ -40,9 +43,14 @@ use datafusion_common::config::TableParquetOptions;
 use datafusion_datasource::TableSchema;
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
+use datafusion_functions::core::file_row_index::FileRowIndexFunc;
+use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::projection::ProjectionExprs;
 use datafusion_physical_expr::{EquivalenceProperties, conjunction};
-use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory;
+use datafusion_physical_expr_adapter::expr_references_scalar_udf;
+use datafusion_physical_expr_adapter::{
+    DefaultPhysicalExprAdapterFactory, rewrite_file_row_index_projection,
+};
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
 use datafusion_physical_plan::DisplayFormatType;
@@ -60,6 +68,7 @@ use datafusion_execution::parquet_encryption::EncryptionFactory;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 use itertools::Itertools;
 use object_store::ObjectStore;
+use parquet::arrow::RowNumber;
 #[cfg(feature = "parquet_encryption")]
 use parquet::encryption::decrypt::FileDecryptionProperties;
 
@@ -669,7 +678,26 @@ impl FileSource for ParquetSource {
         projection: &ProjectionExprs,
     ) -> datafusion_common::Result<Option<Arc<dyn FileSource>>> {
         let mut source = self.clone();
-        source.projection = self.projection.try_merge(projection)?;
+
+        if !projection.iter().any(|projection_expr| {
+            expr_references_scalar_udf::<FileRowIndexFunc>(&projection_expr.expr)
+        }) {
+            source.projection = self.projection.try_merge(projection)?;
+            return Ok(Some(Arc::new(source)));
+        }
+
+        // If we can find a reference to `FileRowIndexFunc`, we add it as a virtual column
+        // or re-use an existing one in the table's schema.
+        let (table_schema, row_index_col) =
+            table_schema_with_row_index_col(self.table_schema());
+
+        source.table_schema = table_schema;
+        source.projection = rewrite_file_row_index_projection(
+            &self.projection,
+            projection,
+            &row_index_col,
+        )?;
+
         Ok(Some(Arc::new(source)))
     }
 
@@ -952,15 +980,12 @@ impl FileSource for ParquetSource {
             reversed_eq_properties.ordering_satisfy(order.iter().cloned())?;
         let sort_order = LexOrdering::new(order.iter().cloned());
         let column_in_file_schema = sort_order.as_ref().is_some_and(|s| {
-            s.first()
-                .expr
-                .downcast_ref::<datafusion_physical_expr::expressions::Column>()
-                .is_some_and(|col| {
-                    self.table_schema
-                        .file_schema()
-                        .field_with_name(col.name())
-                        .is_ok()
-                })
+            s.first().expr.downcast_ref::<Column>().is_some_and(|col| {
+                self.table_schema
+                    .file_schema()
+                    .field_with_name(col.name())
+                    .is_ok()
+            })
         });
 
         if !column_in_file_schema && !reversed_satisfies {
@@ -989,6 +1014,61 @@ impl FileSource for ParquetSource {
     }
 }
 
+fn table_schema_with_row_index_col(table_schema: &TableSchema) -> (TableSchema, Column) {
+    let virtual_offset = table_schema.file_schema().fields().len()
+        + table_schema.table_partition_cols().len();
+    if let Some((idx, field)) =
+        table_schema
+            .virtual_columns()
+            .iter()
+            .enumerate()
+            .find(|(_, field)| {
+                field
+                    .extension_type_name()
+                    .is_some_and(|name| name == RowNumber::NAME)
+            })
+    {
+        return (
+            table_schema.clone(),
+            Column::new(field.name(), virtual_offset + idx),
+        );
+    }
+
+    // The hidden field is shared across all files in this scan, but it must
+    // have a unique table-schema name because later rewrites resolve it by
+    // column name and index.
+    let base_row_index_name = "__datafusion_file_row_index";
+    let mut row_index_name = base_row_index_name.to_string();
+    let mut suffix = 0;
+    while table_schema
+        .table_schema()
+        .field_with_name(&row_index_name)
+        .is_ok()
+    {
+        suffix += 1;
+        row_index_name = format!("{base_row_index_name}_{suffix}");
+    }
+
+    let row_index_table_idx = table_schema.table_schema().fields().len();
+    let row_index_field = Arc::new(
+        Field::new(&row_index_name, DataType::Int64, true).with_extension_type(RowNumber),
+    );
+    (
+        TableSchema::builder(Arc::clone(table_schema.file_schema()))
+            .with_table_partition_cols(table_schema.table_partition_cols().clone())
+            .with_virtual_columns(
+                table_schema
+                    .virtual_columns()
+                    .iter()
+                    .cloned()
+                    .chain([row_index_field])
+                    .collect::<Fields>(),
+            )
+            .build(),
+        Column::new(&row_index_name, row_index_table_idx),
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1622,7 +1702,9 @@ mod tests {
         use datafusion_common::config::ConfigOptions;
         use datafusion_datasource::TableSchema;
         use datafusion_expr::{col, lit as logical_lit};
+        use datafusion_functions::core::expr_fn::file_row_index;
         use datafusion_physical_expr::planner::logical2physical;
+        use datafusion_physical_expr_adapter::rewrite_file_row_index_expr;
         use datafusion_physical_plan::filter_pushdown::PushedDown;
         use parquet::arrow::RowNumber;
 
@@ -1652,13 +1734,20 @@ mod tests {
                 .or(col("value").eq(logical_lit(4i64))),
             full_schema,
         );
+        let (_, row_index_col) = table_schema_with_row_index_col(source.table_schema());
+        let row_index = rewrite_file_row_index_expr(
+            logical2physical(&file_row_index().gt(logical_lit(2i64)), full_schema),
+            row_index_col.name(),
+            row_index_col.index(),
+        )
+        .expect("file_row_index should rewrite to the row_number virtual column");
 
         let config = ConfigOptions::default();
         let prop = source
-            .try_pushdown_filters(vec![pushable, virtual_only, mixed], &config)
+            .try_pushdown_filters(vec![pushable, virtual_only, mixed, row_index], &config)
             .expect("try_pushdown_filters must not error");
 
-        assert_eq!(prop.filters.len(), 3);
+        assert_eq!(prop.filters.len(), 4);
         assert!(
             matches!(prop.filters[0], PushedDown::Yes),
             "file-column filter should be pushable"
@@ -1672,5 +1761,10 @@ mod tests {
             "filter mixing a virtual column with a file column must not be \
              pushed down (row filter would silently drop it)"
         );
+        assert!(
+            matches!(prop.filters[3], PushedDown::No),
+            "file_row_index() rewrites to a virtual column and must not be \
+             pushed down"
+        );
     }
 }

datafusion/functions/src/core/file_row_index.rs

@@ -0,0 +1,96 @@
+// 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.
+
+//! Implementation of the `file_row_index` scalar function.
+
+use arrow::datatypes::DataType;
+use datafusion_common::utils::take_function_args;
+use datafusion_common::{Result, ScalarValue};
+use datafusion_doc::Documentation;
+use datafusion_expr::{
+    ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+    Volatility,
+};
+use datafusion_macros::user_doc;
+
+/// Scalar UDF implementation for `file_row_index()`.
+///
+/// File sources that can expose per-file row indexes rewrite this placeholder
+/// function into a source-provided physical expression. Its fallback
+/// evaluation returns NULL because there is no file context outside a scan.
+#[user_doc(
+    doc_section(label = "Other Functions"),
+    description = r#"Returns the zero-based row offset within the source file
+that produced the current row.
+
+The value is scoped to one file, so rows from different files in the same scan
+can have the same row index. This function is intended to be rewritten at
+file-scan time. 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), this
+function returns NULL.
+"#,
+    syntax_example = "file_row_index()",
+    sql_example = r#"```sql
+SELECT file_row_index() FROM t;
+```"#
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct FileRowIndexFunc {
+    signature: Signature,
+}
+
+impl Default for FileRowIndexFunc {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl FileRowIndexFunc {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::nullary(Volatility::Volatile),
+        }
+    }
+}
+
+impl ScalarUDFImpl for FileRowIndexFunc {
+    fn name(&self) -> &str {
+        "file_row_index"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, args: &[DataType]) -> Result<DataType> {
+        let [] = take_function_args(self.name(), args)?;
+        Ok(DataType::Int64)
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let [] = take_function_args(self.name(), args.args)?;
+        Ok(ColumnarValue::Scalar(ScalarValue::Int64(None)))
+    }
+
+    fn placement(&self, _args: &[ExpressionPlacement]) -> ExpressionPlacement {
+        ExpressionPlacement::MoveTowardsLeafNodes
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.doc()
+    }
+}

datafusion/functions/src/core/mod.rs

@@ -28,6 +28,7 @@ pub mod arrowtypeof;
 pub mod cast_to_type;
 pub mod coalesce;
 pub mod expr_ext;
+pub mod file_row_index;
 pub mod getfield;
 pub mod greatest;
 mod greatest_least_utils;
@@ -67,6 +68,7 @@ make_udf_function!(version::VersionFunc, version);
 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);
 
 pub mod expr_fn {
     use datafusion_expr::{Expr, Literal};
@@ -143,6 +145,9 @@ pub mod expr_fn {
         union_tag,
         "Returns the name of the currently selected field in the union",
         arg1
+    ),(
+        file_row_index,
+        "Returns the offset of the row within its source file",
     ));
 
     #[doc = "Returns the value of the field with the given name from the struct"]
@@ -196,5 +201,6 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
         union_tag(),
         version(),
         r#struct(),
+        file_row_index(),
     ]
 }

datafusion/physical-expr-adapter/src/lib.rs

@@ -29,5 +29,6 @@ pub mod schema_rewriter;
 pub use schema_rewriter::{
     BatchAdapter, BatchAdapterFactory, DefaultPhysicalExprAdapter,
     DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory,
-    replace_columns_with_literals,
+    expr_references_scalar_udf, replace_columns_with_literals,
+    rewrite_file_row_index_expr, rewrite_file_row_index_projection, rewrite_scalar_udf,
 };