Fix non-deterministic iteration in SessionStateBuilder

Author: shehab-ali

datafusion/core/src/execution/session_state.rs

@@ -978,6 +978,20 @@ impl SessionState {
 /// A builder to be used for building [`SessionState`]'s. Defaults will
 /// be used for all values unless explicitly provided.
 ///
+/// Deduplicates function-registry map entries by keeping only entries whose key
+/// matches the canonical name. The session stores one hash map entry per alias
+/// plus the canonical name; filtering to canonical-name entries yields exactly
+/// one [`Arc`] per logical function.
+fn dedup_function_registry_by_canonical_name<T>(
+    map: &HashMap<String, Arc<T>>,
+    canonical_name: impl Fn(&T) -> &str,
+) -> Vec<Arc<T>> {
+    map.iter()
+        .filter(|(key, udf)| key.as_str() == canonical_name(udf.as_ref()))
+        .map(|(_, udf)| Arc::clone(udf))
+        .collect()
+}
+
 /// See example on [`SessionState`]
 #[derive(Clone)]
 pub struct SessionStateBuilder {
@@ -1088,11 +1102,18 @@ impl SessionStateBuilder {
             query_planner: Some(existing.query_planner),
             catalog_list: Some(existing.catalog_list),
             table_functions: Some(existing.table_functions),
-            scalar_functions: Some(existing.scalar_functions.into_values().collect_vec()),
-            aggregate_functions: Some(
-                existing.aggregate_functions.into_values().collect_vec(),
-            ),
-            window_functions: Some(existing.window_functions.into_values().collect_vec()),
+            scalar_functions: Some(dedup_function_registry_by_canonical_name(
+                &existing.scalar_functions,
+                |u| u.name(),
+            )),
+            aggregate_functions: Some(dedup_function_registry_by_canonical_name(
+                &existing.aggregate_functions,
+                |u| u.name(),
+            )),
+            window_functions: Some(dedup_function_registry_by_canonical_name(
+                &existing.window_functions,
+                |u| u.name(),
+            )),
             extension_types: Some(existing.extension_types),
             serializer_registry: Some(existing.serializer_registry),
             file_formats: Some(existing.file_formats.into_values().collect_vec()),
@@ -1862,8 +1883,6 @@ impl ContextProvider for SessionContextProvider<'_> {
         name: &str,
         args: Vec<Expr>,
     ) -> datafusion_common::Result<Arc<dyn TableSource>> {
-        use datafusion_catalog::TableFunctionArgs;
-
         let tbl_func = self
             .state
             .table_functions
@@ -1886,8 +1905,7 @@ impl ContextProvider for SessionContextProvider<'_> {
                     .and_then(|e| simplifier.simplify(e))
             })
             .collect::<datafusion_common::Result<Vec<_>>>()?;
-        let provider = tbl_func
-            .create_table_provider_with_args(TableFunctionArgs::new(&args, self.state))?;
+        let provider = tbl_func.create_table_provider(&args)?;
 
         Ok(provider_as_source(provider))
     }
@@ -2340,6 +2358,37 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn new_from_existing_preserves_scalar_udf_aliases() -> Result<()> {
+        use arrow::datatypes::DataType;
+        use datafusion_common::ScalarValue;
+        use datafusion_expr::registry::FunctionRegistry;
+        use datafusion_expr::{ColumnarValue, Volatility, create_udf};
+
+        let udf = create_udf(
+            "postgres_to_char",
+            vec![DataType::Utf8],
+            DataType::Utf8,
+            Volatility::Immutable,
+            Arc::new(|_args| Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))),
+        )
+        .with_aliases(["to_char"]);
+
+        let mut state = SessionStateBuilder::new().build();
+        state.register_udf(Arc::new(udf))?;
+
+        assert_eq!(state.udf("postgres_to_char")?.name(), "postgres_to_char");
+        assert_eq!(state.udf("to_char")?.name(), "postgres_to_char");
+
+        let roundtrip = SessionStateBuilder::new_from_existing(state).build();
+        assert_eq!(roundtrip.udf("to_char")?.name(), "postgres_to_char");
+        assert_eq!(
+            roundtrip.udf("postgres_to_char")?.name(),
+            "postgres_to_char"
+        );
+        Ok(())
+    }
+
     #[test]
     fn test_session_state_with_optimizer_rules() {
         #[derive(Default, Debug)]