Skip to content

fix: use qualified function identifiers for temp functions #900

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

Merged
SemyonSinchenko merged 1 commit into from
Sep 8, 2026
Merged

fix: use qualified function identifiers for temp functions #900

Show file tree
Hide file tree
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
10 changes: 6 additions & 4 deletions core/src/main/scala/org/graphframes/lib/KCore.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,14 +65,15 @@ object KCore extends Serializable with Logging {
storageLevel: StorageLevel,
checkpointInterval: Int,
useLocalCheckpoints: Boolean): DataFrame = {
val spark = graph.vertices.sparkSession
val degrees = graph.degrees
val preparedGraph = GraphFrame(
degrees.withColumn("degree", col("degree").cast(IntegerType)),
graph.edges.select(GraphFrame.SRC, GraphFrame.DST))

val functionRegistry = graph.vertices.sparkSession.sessionState.functionRegistry
functionRegistry.registerFunction(
FunctionIdentifier("_kcoreMerge"),
val functionRegistry = spark.sessionState.functionRegistry
functionRegistry.createOrReplaceTempFunction(
"_kcoreMerge",
(children: Seq[Expression]) => KCoreMerge(children(0), children(1)),
"scala_udf")

Expand All @@ -98,7 +99,8 @@ object KCore extends Serializable with Logging {

pregel.run()
} finally {
val dereg = functionRegistry.dropFunction(FunctionIdentifier("_kcoreMerge"))
val catalog = spark.sessionState.catalog
val dereg = catalog.unregisterFunction(FunctionIdentifier("_kcoreMerge"))
if (!dereg) {
logWarn(
"graphframes faced an internal error and was not able to de-register function _kcoreMerge; Spark' functionRegistry is in a bad state")
Expand Down
7 changes: 4 additions & 3 deletions core/src/main/scala/org/graphframes/lib/RandomizedContraction.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,8 @@ private[graphframes] object RandomizedContraction extends Logging with Serializa
logInfo(s"$logPrefix Using $checkpointDir for storing intermediate tables.")

val functionRegistry = spark.sessionState.functionRegistry
functionRegistry.registerFunction(
FunctionIdentifier("_axpb"),
functionRegistry.createOrReplaceTempFunction(
"_axpb",
(children: Seq[Expression]) => FiniteAXPlusB(children(0), children(1), children(2)),
"scala_udf")

Expand Down Expand Up @@ -302,7 +302,8 @@ private[graphframes] object RandomizedContraction extends Logging with Serializa
// to be 100% sure;
edges.unpersist()
cleanupCheckpointDir()
val dereg = functionRegistry.dropFunction(FunctionIdentifier("_axpb"))
val catalog = spark.sessionState.catalog
val dereg = catalog.unregisterFunction(FunctionIdentifier("_axpb"))
if (!dereg) {
logWarn(
"graphframes faced an internal error and was not able to de-register function _axpb; Spark' functionRegistry is in a bad state")
Expand Down
3 changes: 1 addition & 2 deletions core/src/test/scala/org/graphframes/lib/RandomizedContractionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,8 +280,7 @@ class RandomizedContractionSuite extends SparkFunSuite with GraphFrameTestSparkC
}

private def assertFunctionRegistryClean(): Unit = {
val functionRegistry = spark.sessionState.functionRegistry
val _ = assert(!functionRegistry.functionExists(FunctionIdentifier("_axpb")))
val _ = assert(!spark.sessionState.catalog.isTemporaryFunction(FunctionIdentifier("_axpb")))
}

private def listParquetFiles(): Set[String] = {
Expand Down
Loading