use arrow::datatypes::DataType;
use datafusion_common::utils::take_function_args;
use datafusion_common::{Result, exec_err};
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 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), direct
evaluation returns an error.
"#,
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)?;
exec_err!("file_row_index() is source dependent and cannot be evaluated directly")
}
fn placement(&self, _args: &[ExpressionPlacement]) -> ExpressionPlacement {
ExpressionPlacement::MoveTowardsLeafNodes
}
fn documentation(&self) -> Option<&Documentation> {
self.doc()
}
}