Skip to content

Fix non-deterministic iteration in SessionStateBuilder #21262

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shehab-ali wants to merge 1 commit into
base: main
Choose a base branch
from
+57 −5
Open

Fix non-deterministic iteration in SessionStateBuilder #21262

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 57 additions & 5 deletions datafusion/core/src/execution/session_state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -975,6 +975,20 @@ impl SessionState {
}
}

/// Deduplicates function-registry map entries by keeping only entries whose key
Comment thread
shehab-ali marked this conversation as resolved.
/// 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>(
Copy link
Copy Markdown
Contributor

@ahmed-mez ahmed-mez Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: I think the function could consume the map instead of borrowing it, to avoid some arc clones

  fn dedup_function_registry_by_canonical_name<T>(
      map: HashMap<String, Arc<T>>,
      canonical_name: impl Fn(&T) -> &str,
  ) -> Vec<Arc<T>> {
      map.into_iter()
          .filter(|(key, udf)| key.as_str() == canonical_name(udf.as_ref()))
          .map(|(_, udf)| udf)  // no Arc::clone needed
          .collect()
  }

Copy link
Copy Markdown
Contributor

@ahmed-mez ahmed-mez Apr 8, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alternatively, it might be more robust to dedup by identity

  fn dedup_by_identity<T>(map: HashMap<String, Arc<T>>) -> Vec<Arc<T>> {
      let mut seen = HashSet::new();
      map.into_values()
          .filter(|arc| seen.insert(Arc::as_ptr(arc)))
          .collect()
  }

map: &HashMap<String, Arc<T>>,
canonical_name: impl Fn(&T) -> &str,
Copy link
Copy Markdown
Contributor

@alexanderbianchi alexanderbianchi Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need canonical_name as it's own function if all the uses are the same function?

Copy link
Copy Markdown
Author

@shehab-ali shehab-ali Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

those function have different types though: WindowUDF, ScalarUDF, AggregateUDF

Copy link
Copy Markdown
Contributor

@alexanderbianchi alexanderbianchi Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah I see now, my bad.

) -> Vec<Arc<T>> {
map.iter()
.filter(|(key, udf)| key.as_str() == canonical_name(udf.as_ref()))
.map(|(_, udf)| Arc::clone(udf))
.collect()
}

/// A builder to be used for building [`SessionState`]'s. Defaults will
/// be used for all values unless explicitly provided.
///
Expand Down Expand Up @@ -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()),
Expand Down Expand Up @@ -2340,6 +2361,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)]
Expand Down
Loading