Support volatile scalar functions (#639)

Based on #636 by @tobilg, but more idiomatic using a new trait method.

Author: mlafeldt

crates/duckdb/src/vscalar/arrow.rs

@@ -84,6 +84,21 @@ pub trait VArrowScalar: Sized {
     /// The possible signatures of the scalar function. These will result in DuckDB scalar function overloads.
     /// The invoke method should be able to handle all of these signatures.
     fn signatures() -> Vec<ArrowFunctionSignature>;
+
+    /// Whether the scalar function is volatile.
+    ///
+    /// Volatile functions are re-evaluated for each row, even if they have no parameters.
+    /// This is useful for functions that generate random or unique values, such as random
+    /// number generators, UUID generators, or fake data generators.
+    ///
+    /// By default, DuckDB optimizes zero-argument scalar functions as constants, evaluating
+    /// them only once. Returning true from this method prevents this optimization.
+    ///
+    /// # Default
+    /// Returns `false` by default, meaning the function is not volatile.
+    fn volatile() -> bool {
+        false
+    }
 }
 
 impl<T> VScalar for T

crates/duckdb/src/vscalar/function.rs

@@ -53,7 +53,8 @@ use libduckdb_sys::{
     duckdb_create_scalar_function_set, duckdb_data_chunk, duckdb_delete_callback_t, duckdb_destroy_scalar_function,
     duckdb_function_info, duckdb_scalar_function, duckdb_scalar_function_add_parameter, duckdb_scalar_function_set,
     duckdb_scalar_function_set_extra_info, duckdb_scalar_function_set_function, duckdb_scalar_function_set_name,
-    duckdb_scalar_function_set_return_type, duckdb_scalar_function_set_varargs, duckdb_vector, DuckDBSuccess,
+    duckdb_scalar_function_set_return_type, duckdb_scalar_function_set_varargs, duckdb_scalar_function_set_volatile,
+    duckdb_vector, DuckDBSuccess,
 };
 
 use crate::{core::LogicalTypeHandle, Error};
@@ -112,6 +113,21 @@ impl ScalarFunction {
         self
     }
 
+    /// Marks the scalar function as volatile.
+    ///
+    /// Volatile functions are re-evaluated for each row, even if they have no parameters.
+    /// This is useful for functions that generate random or unique values, such as random
+    /// number generators, UUID generators, or fake data generators.
+    ///
+    /// By default, DuckDB optimizes zero-argument scalar functions as constants, evaluating
+    /// them only once. Setting a function as volatile prevents this optimization.
+    pub fn set_volatile(&self) -> &Self {
+        unsafe {
+            duckdb_scalar_function_set_volatile(self.ptr);
+        }
+        self
+    }
+
     /// Assigns extra information to the scalar function using raw pointers.
     ///
     /// For most use cases, prefer [`set_extra_info`](Self::set_extra_info) which handles memory management automatically.

crates/duckdb/src/vscalar/mod.rs

@@ -45,6 +45,21 @@ pub trait VScalar: Sized {
     /// These will result in DuckDB scalar function overloads.
     /// The invoke method should be able to handle all of these signatures.
     fn signatures() -> Vec<ScalarFunctionSignature>;
+
+    /// Whether the scalar function is volatile.
+    ///
+    /// Volatile functions are re-evaluated for each row, even if they have no parameters.
+    /// This is useful for functions that generate random or unique values, such as random
+    /// number generators, UUID generators, or fake data generators.
+    ///
+    /// By default, DuckDB optimizes zero-argument scalar functions as constants, evaluating
+    /// them only once. Returning true from this method prevents this optimization.
+    ///
+    /// # Default
+    /// Returns `false` by default, meaning the function is not volatile.
+    fn volatile() -> bool {
+        false
+    }
 }
 
 /// Duckdb scalar function parameters
@@ -144,6 +159,9 @@ impl Connection {
             let scalar_function = ScalarFunction::new(name)?;
             signature.register_with_scalar(&scalar_function);
             scalar_function.set_function(Some(scalar_func::<S>));
+            if S::volatile() {
+                scalar_function.set_volatile();
+            }
             scalar_function.set_extra_info(S::State::default());
             set.add_function(scalar_function)?;
         }
@@ -163,6 +181,9 @@ impl Connection {
             let scalar_function = ScalarFunction::new(name)?;
             signature.register_with_scalar(&scalar_function);
             scalar_function.set_function(Some(scalar_func::<S>));
+            if S::volatile() {
+                scalar_function.set_volatile();
+            }
             scalar_function.set_extra_info(state.clone());
             set.add_function(scalar_function)?;
         }
@@ -374,4 +395,103 @@ mod test {
 
         Ok(())
     }
+
+    // Counters for testing volatile functions
+    use std::sync::atomic::{AtomicU64, Ordering};
+    static VOLATILE_COUNTER: AtomicU64 = AtomicU64::new(0);
+    static NON_VOLATILE_COUNTER: AtomicU64 = AtomicU64::new(0);
+
+    struct CounterScalar {}
+
+    impl VScalar for CounterScalar {
+        type State = ();
+
+        unsafe fn invoke(
+            _: &Self::State,
+            input: &mut DataChunkHandle,
+            output: &mut dyn WritableVector,
+        ) -> Result<(), Box<dyn std::error::Error>> {
+            let len = input.len();
+            let mut output_vec = output.flat_vector();
+            let data = output_vec.as_mut_slice::<i64>();
+
+            for item in data.iter_mut().take(len) {
+                *item = NON_VOLATILE_COUNTER.fetch_add(1, Ordering::SeqCst) as i64;
+            }
+            Ok(())
+        }
+
+        fn signatures() -> Vec<ScalarFunctionSignature> {
+            vec![ScalarFunctionSignature::exact(
+                vec![],
+                LogicalTypeHandle::from(LogicalTypeId::Bigint),
+            )]
+        }
+    }
+
+    struct VolatileCounterScalar {}
+
+    impl VScalar for VolatileCounterScalar {
+        type State = ();
+
+        unsafe fn invoke(
+            _: &Self::State,
+            input: &mut DataChunkHandle,
+            output: &mut dyn WritableVector,
+        ) -> Result<(), Box<dyn std::error::Error>> {
+            let len = input.len();
+            let mut output_vec = output.flat_vector();
+            let data = output_vec.as_mut_slice::<i64>();
+
+            for item in data.iter_mut().take(len) {
+                *item = VOLATILE_COUNTER.fetch_add(1, Ordering::SeqCst) as i64;
+            }
+            Ok(())
+        }
+
+        fn signatures() -> Vec<ScalarFunctionSignature> {
+            vec![ScalarFunctionSignature::exact(
+                vec![],
+                LogicalTypeHandle::from(LogicalTypeId::Bigint),
+            )]
+        }
+
+        fn volatile() -> bool {
+            true
+        }
+    }
+
+    #[test]
+    fn test_volatile_scalar() -> Result<(), Box<dyn Error>> {
+        let conn = Connection::open_in_memory()?;
+
+        VOLATILE_COUNTER.store(0, Ordering::SeqCst);
+        conn.register_scalar_function::<VolatileCounterScalar>("volatile_counter")?;
+
+        let values: Vec<i64> = conn
+            .prepare("SELECT volatile_counter() FROM generate_series(1, 5)")?
+            .query_map([], |row| row.get(0))?
+            .collect::<Result<_, _>>()?;
+
+        assert_eq!(values, [0, 1, 2, 3, 4]);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_non_volatile_scalar() -> Result<(), Box<dyn Error>> {
+        let conn = Connection::open_in_memory()?;
+
+        NON_VOLATILE_COUNTER.store(0, Ordering::SeqCst);
+        conn.register_scalar_function::<CounterScalar>("non_volatile_counter")?;
+
+        // Constant folding should make every row identical
+        let distinct_count: i64 = conn
+            .prepare("SELECT COUNT(DISTINCT non_volatile_counter()) FROM generate_series(1, 5)")?
+            .query_row([], |row| row.get(0))?;
+
+        assert_eq!(distinct_count, 1);
+
+        Ok(())
+    }
 }