-
Notifications
You must be signed in to change notification settings - Fork 140
feat(workflow-operator): add Python UDF UI parameter injection model #5141
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
4ec7c7a
f386522
4652aac
962a52c
b18d708
16373a4
75fe048
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,197 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
| package org.apache.texera.amber.operator.udf.python | ||
|
|
||
| import org.apache.texera.amber.core.tuple.{Attribute, AttributeType} | ||
| import org.apache.texera.amber.pybuilder.PythonTemplateBuilder | ||
| import org.apache.texera.amber.pybuilder.PythonTemplateBuilder.PythonTemplateBuilderStringContext | ||
|
|
||
| import scala.util.matching.Regex | ||
|
|
||
| /** | ||
| * Injects the reserved UI-parameter hook into user-written Python UDF code. | ||
| * | ||
| * Operator descriptors should call this after loading saved [[UiUDFParameter]] values and before sending Python source | ||
| * to runtime execution. The injected hook returns decoded parameter names and values that Python runtime support reads | ||
| * before the user's `open()` method runs. | ||
| */ | ||
| object PythonUdfUiParameterInjector { | ||
|
|
||
| private val InjectedUiParametersHookMethodName = "_texera_injected_ui_parameters" | ||
| private val InjectedUiParametersHookMethodHeader = | ||
| s"def $InjectedUiParametersHookMethodName(self) -> Dict[str, Any]:" | ||
| private val UnsupportedUiParameterTypes = Set(AttributeType.BINARY, AttributeType.LARGE_BINARY) | ||
|
|
||
| // Keep supported user-facing UDF class names in sync with the frontend parser. | ||
| private val SupportedPythonUdfClassHeaderRegex: Regex = | ||
| """(?m)^([ \t]*)class\s+(ProcessTupleOperator|ProcessBatchOperator|ProcessTableOperator|GenerateOperator)\s*\([^)]*\)\s*:\s*(?:#.*)?$""".r | ||
|
|
||
| private def validate(uiParameters: List[UiUDFParameter]): Unit = { | ||
| val attributes = uiParameters.map(parameterAttribute) | ||
| attributes.foreach(validateSupportedType) | ||
|
|
||
| attributes | ||
| .groupBy(_.getName) | ||
| .collectFirst { | ||
| case (parameterName, matchingAttributes) if matchingAttributes.size > 1 => parameterName | ||
| } | ||
| .foreach { duplicateName => | ||
| throw new RuntimeException(s"UiParameter name '$duplicateName' is declared more than once.") | ||
| } | ||
| } | ||
|
|
||
| private def parameterAttribute(parameter: UiUDFParameter): Attribute = | ||
| Option(parameter).flatMap(parameter => Option(parameter.attribute)).getOrElse { | ||
| throw new RuntimeException("UiParameter attribute is required.") | ||
| } | ||
|
|
||
| private def validateSupportedType(attribute: Attribute): Unit = { | ||
| if (UnsupportedUiParameterTypes.contains(attribute.getType)) { | ||
| throw new RuntimeException( | ||
| s"UiParameter type '${attribute.getType.name()}' is not supported. " + | ||
| "Use string, integer, long, double, boolean, or timestamp instead." | ||
| ) | ||
| } | ||
| } | ||
|
|
||
| private def buildInjectedParameterEntry(parameter: UiUDFParameter): PythonTemplateBuilder = { | ||
| pyb"${parameter.attribute.getName}: ${parameter.value}" | ||
| } | ||
|
|
||
| private def buildInjectedParametersMap( | ||
| uiParameters: List[UiUDFParameter] | ||
| ): PythonTemplateBuilder = { | ||
| val entries = uiParameters.map(buildInjectedParameterEntry) | ||
| entries.reduceOption((acc, entry) => acc + pyb", " + entry).getOrElse(pyb"") | ||
| } | ||
|
|
||
| private def buildInjectedHookMethod(uiParameters: List[UiUDFParameter]): String = { | ||
| val injectedParametersMap = buildInjectedParametersMap(uiParameters) | ||
|
|
||
| (pyb"""|@overrides | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There's no |
||
| |$InjectedUiParametersHookMethodHeader | ||
| | return {""" + | ||
| injectedParametersMap + | ||
| pyb"""} | ||
| |""").encode | ||
| } | ||
|
|
||
| private def indentBlock(block: String, indent: String): String = { | ||
| block | ||
| .split("\n", -1) | ||
| .map { line => | ||
| if (line.nonEmpty) indent + line else line | ||
| } | ||
| .mkString("\n") | ||
| } | ||
|
|
||
| private def lineEndIndex(text: String, from: Int): Int = { | ||
| val lineEnd = text.indexOf('\n', from) | ||
| if (lineEnd < 0) text.length else lineEnd | ||
| } | ||
|
|
||
| private def detectClassBlockEnd(code: String, classHeaderStart: Int, classIndent: String): Int = { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This decides the class body ends at the first non-blank line indented at or below the class header. That's right for ordinary code, but it doesn't account for string literals, so a triple-quoted block inside a method whose contents start at column 0 looks like the end of the class — and we'd splice the hook into the middle of the string: class ProcessTupleOperator(UDFOperatorV2):
def process_tuple(self, tuple_, port):
sql = """
SELECT * FROM t
"""
yield tuple_
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I can take a look about this one, thanks |
||
| val classLineEnd = lineEndIndex(code, classHeaderStart) | ||
| var lineStart = if (classLineEnd < code.length) classLineEnd + 1 else code.length | ||
|
|
||
| while (lineStart < code.length) { | ||
| val lineEnd = lineEndIndex(code, lineStart) | ||
| val line = code.substring(lineStart, lineEnd) | ||
|
|
||
| val trimmed = line.trim | ||
| val isBlank = trimmed.isEmpty | ||
|
|
||
| val currentIndentLen = line.segmentLength(ch => ch == ' ' || ch == '\t') | ||
| val classIndentLen = classIndent.length | ||
|
|
||
| if (!isBlank && currentIndentLen <= classIndentLen) { | ||
| return lineStart | ||
| } | ||
|
|
||
| lineStart = if (lineEnd < code.length) lineEnd + 1 else code.length | ||
| } | ||
|
|
||
| code.length | ||
| } | ||
|
|
||
| private def containsReservedHook(classBlock: String): Boolean = { | ||
| val hookRegex = | ||
| ("""(?m)^[ \t]+def\s+""" + Regex.quote(InjectedUiParametersHookMethodName) + """\s*\(""").r | ||
| hookRegex.findFirstIn(classBlock).isDefined | ||
| } | ||
|
|
||
| private def injectHookIntoUserClass(encodedUserCode: String, hookMethod: String): String = { | ||
| val classHeaderMatch = | ||
| SupportedPythonUdfClassHeaderRegex.findFirstMatchIn(encodedUserCode).getOrElse { | ||
| return encodedUserCode | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. When the code doesn't contain one of the four supported class names, we hand it back unchanged even if there are parameters to inject. So if a user renames their class, their configured parameters just disappear at runtime with no signal. Since
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Got it, I will update it |
||
| } | ||
|
|
||
| val classHeaderStart = classHeaderMatch.start | ||
| val classIndent = classHeaderMatch.group(1) | ||
| val classBlockEnd = detectClassBlockEnd(encodedUserCode, classHeaderStart, classIndent) | ||
|
|
||
| val classBlock = encodedUserCode.substring(classHeaderStart, classBlockEnd) | ||
|
|
||
| if (containsReservedHook(classBlock)) { | ||
| throw new RuntimeException( | ||
| s"Reserved method '$InjectedUiParametersHookMethodName' is already defined in the UDF class. Please rename your method." | ||
| ) | ||
| } | ||
|
|
||
| val bodyIndent = inferClassBodyIndent(classBlock, classIndent).getOrElse(classIndent + " ") | ||
| val indentedHook = indentBlock( | ||
| (if (classBlock.endsWith("\n")) "" else "\n") + hookMethod.trim + "\n", | ||
| bodyIndent | ||
| ) | ||
|
|
||
| encodedUserCode.substring(0, classBlockEnd) + | ||
| indentedHook + | ||
| encodedUserCode.substring(classBlockEnd) | ||
| } | ||
|
|
||
| private def inferClassBodyIndent(classBlock: String, classIndent: String): Option[String] = { | ||
| val lines = classBlock.split("\n", -1).toList.drop(1) | ||
|
|
||
| lines.collectFirst { | ||
| case line if line.trim.nonEmpty => | ||
| val leading = line.takeWhile(ch => ch == ' ' || ch == '\t') | ||
| if (leading.length > classIndent.length) leading else classIndent + " " | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Returns Python code with the UI-parameter hook injected into the supported UDF class. | ||
| * | ||
| * If `uiParameters` is empty, the code is only passed through normal Python-template encoding. Throws | ||
| * [[RuntimeException]] when parameter metadata is invalid or the user already defines the reserved hook method. | ||
| */ | ||
| def inject(code: String, uiParameters: List[UiUDFParameter]): String = { | ||
| val parameters = Option(uiParameters).getOrElse(List.empty) | ||
| validate(parameters) | ||
|
|
||
| val encodedUserCode = pyb"$code".encode | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Minor: running the whole program through
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'll take care of it. |
||
|
|
||
| if (parameters.isEmpty) { | ||
| return encodedUserCode | ||
| } | ||
|
|
||
| val hookMethod = buildInjectedHookMethod(parameters) | ||
| injectHookIntoUserClass(encodedUserCode, hookMethod) | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,47 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
|
|
||
| package org.apache.texera.amber.operator.udf.python | ||
|
|
||
| import com.fasterxml.jackson.annotation.JsonProperty | ||
| import com.kjetland.jackson.jsonSchema.annotations.JsonSchemaTitle | ||
| import org.apache.texera.amber.core.tuple.Attribute | ||
| import org.apache.texera.amber.pybuilder.PyStringTypes.EncodableString | ||
|
|
||
| import javax.validation.Valid | ||
| import javax.validation.constraints.NotNull | ||
|
|
||
| /** | ||
| * Serialized operator property for one Python UDF UI parameter. | ||
| * | ||
| * `attribute` carries the inferred parameter name and type. `value` is user-entered text and is marked as | ||
| * [[EncodableString]] so Python code generation decodes it at runtime instead of embedding raw text into generated code. | ||
| */ | ||
| class UiUDFParameter { | ||
|
|
||
| @JsonProperty(required = true) | ||
| @JsonSchemaTitle("Attribute") | ||
| @Valid | ||
| @NotNull(message = "Attribute is required") | ||
| var attribute: Attribute = _ | ||
|
|
||
| @JsonProperty() | ||
| @JsonSchemaTitle("Value") | ||
| var value: EncodableString = "" | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
DictandAnyaren't actually exported byfrom pytexera import *— its__all__only re-exportsIterator,Optional, andUnionfrom typing. So in the user's module these names are undefined, and since we run Python 3.10–3.13 (where return annotations are evaluated at definition time), the class raisesNameError: name 'Dict' is not definedthe moment it loads. The test that asserts the output shouldn't containimport typingquietly locks this assumption in, so it'll stay green while generating code that can't run. We could addDict/Anytopytexera's__all__, inject afrom typing import Dict, Any, or just drop the return annotation — whichever we pick, it'd be good to settle it before the runtime PR starts depending on it.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is covered in next PR actually but I did not include it here since this is for scala changes. Next PR will be about Python side.