diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/java/JavaUDFOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/java/JavaUDFOpDescSpec.scala index 6ea1fe8bc47..586ab99c5f3 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/java/JavaUDFOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/java/JavaUDFOpDescSpec.scala @@ -22,7 +22,8 @@ package org.apache.texera.amber.operator.udf.java import org.apache.texera.amber.core.executor.OpExecWithCode import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} -import org.apache.texera.amber.operator.LogicalOp +import org.apache.texera.amber.core.workflow.UnknownPartition +import org.apache.texera.amber.operator.{LogicalOp, PortDescription} import org.apache.texera.amber.operator.metadata.OperatorGroupConstants import org.apache.texera.amber.util.JSONUtils.objectMapper import org.scalatest.flatspec.AnyFlatSpec @@ -106,4 +107,75 @@ class JavaUDFOpDescSpec extends AnyFlatSpec with Matchers { j.workers shouldBe 4 j.retainInputColumns shouldBe true } + + "JavaUDFOpDesc.getPhysicalOp" should + "build a parallelizable one-to-many op when workers > 1" in { + val d = new JavaUDFOpDesc + d.code = "return t;" + d.workers = 2 + val physical = d.getPhysicalOp(workflowId, executionId) + physical.parallelizable shouldBe true + physical.isOneToManyOp shouldBe true + physical.suggestedWorkerNum shouldBe Some(2) + physical.opExecInitInfo match { + case OpExecWithCode(code, language) => + code shouldBe "return t;" + language shouldBe "java" + case other => fail(s"expected OpExecWithCode, got $other") + } + } + + "JavaUDFOpDesc.operatorInfo" should + "derive input ports from a configured inputPorts list" in { + val d = new JavaUDFOpDesc + d.inputPorts = List( + PortDescription( + portID = "0", + displayName = "left", + disallowMultiInputs = true, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ) + ) + val info = d.operatorInfo + info.inputPorts should have length 1 + info.inputPorts.head.displayName shouldBe "left" + // also drives the getPhysicalOp inputPorts != null partitionRequirement branch + val physical = d.getPhysicalOp(workflowId, executionId) + physical.inputPorts.keySet shouldBe info.inputPorts.map(_.id).toSet + physical.partitionRequirement shouldBe List(Some(UnknownPartition())) + } + + it should "derive output ports from a configured outputPorts list" in { + val d = new JavaUDFOpDesc + d.outputPorts = List( + PortDescription( + portID = "0", + displayName = "result", + disallowMultiInputs = false, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ) + ) + val info = d.operatorInfo + info.outputPorts should have length 1 + info.outputPorts.head.displayName shouldBe "result" + } + + "JavaUDFOpDesc.runtimeReconfiguration" should + "return the new op's physical op with no state transfer" in { + val oldOp = new JavaUDFOpDesc + val newOp = new JavaUDFOpDesc + newOp.code = "return t2;" + val result = oldOp.runtimeReconfiguration(workflowId, executionId, oldOp, newOp) + result.isSuccess shouldBe true + val (physical, stateTransfer) = result.get + stateTransfer shouldBe None + physical.opExecInitInfo match { + case OpExecWithCode(code, _) => code shouldBe "return t2;" + case other => fail(s"expected OpExecWithCode, got $other") + } + } } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala index be319f4e411..aea38193aa9 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/python/PythonUDFOpDescV2Spec.scala @@ -22,7 +22,8 @@ package org.apache.texera.amber.operator.udf.python import org.apache.texera.amber.core.executor.OpExecWithCode import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} -import org.apache.texera.amber.operator.LogicalOp +import org.apache.texera.amber.core.workflow.{PortIdentity, UnknownPartition} +import org.apache.texera.amber.operator.{LogicalOp, PortDescription} import org.apache.texera.amber.operator.metadata.OperatorGroupConstants import org.apache.texera.amber.util.JSONUtils.objectMapper import org.scalatest.flatspec.AnyFlatSpec @@ -134,4 +135,106 @@ class PythonUDFOpDescV2Spec extends AnyFlatSpec with Matchers { p.envName shouldBe "myenv" p.outputColumns shouldBe List(new Attribute("res", AttributeType.INTEGER)) } + + "PythonUDFOpDescV2.getPhysicalOp" should + "use a parallelizable one-to-one op with the suggested worker count when workers > 1" in { + val d = new PythonUDFOpDescV2 + d.code = "yield t" + d.workers = 4 + val physical = d.getPhysicalOp(workflowId, executionId) + physical.parallelizable shouldBe true + physical.suggestedWorkerNum shouldBe Some(4) + physical.opExecInitInfo match { + case OpExecWithCode(code, language) => + code shouldBe "yield t" + language shouldBe "python" + case other => fail(s"expected OpExecWithCode, got $other") + } + } + + it should "carry the trimmed virtual-environment name when the default env is disabled" in { + val d = new PythonUDFOpDescV2 + d.defaultEnv = false + d.envName = " myenv " + d.getPhysicalOp(workflowId, executionId).pveName shouldBe "myenv" + } + + it should "map each custom input port's partitionRequirement into the physical op" in { + val d = new PythonUDFOpDescV2 + d.code = "yield t" + d.inputPorts = List( + PortDescription( + portID = "p0", + displayName = "in", + disallowMultiInputs = false, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ) + ) + d.getPhysicalOp(workflowId, executionId).partitionRequirement shouldBe List( + Some(UnknownPartition()) + ) + } + + "PythonUDFOpDescV2.operatorInfo" should + "derive input ports from a custom inputPorts descriptor list" in { + val d = new PythonUDFOpDescV2 + d.inputPorts = List( + PortDescription( + portID = "p0", + displayName = "left", + disallowMultiInputs = true, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ), + PortDescription( + portID = "p1", + displayName = "right", + disallowMultiInputs = false, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List(0) + ) + ) + val info = d.operatorInfo + info.inputPorts should have length 2 + info.inputPorts.head.displayName shouldBe "left" + info.inputPorts.head.disallowMultiLinks shouldBe true + info.inputPorts(1).displayName shouldBe "right" + info.inputPorts(1).dependencies shouldBe List(PortIdentity(0)) + } + + it should "derive output ports from a custom outputPorts descriptor list" in { + val d = new PythonUDFOpDescV2 + d.outputPorts = List( + PortDescription( + portID = "o0", + displayName = "result", + disallowMultiInputs = false, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ) + ) + val info = d.operatorInfo + info.outputPorts should have length 1 + info.outputPorts.head.displayName shouldBe "result" + } + + "PythonUDFOpDescV2.runtimeReconfiguration" should + "return the new op's physical op with no state transfer" in { + val oldOp = new PythonUDFOpDescV2 + val newOp = new PythonUDFOpDescV2 + newOp.code = "yield reconfigured" + val result = oldOp.runtimeReconfiguration(workflowId, executionId, oldOp, newOp) + result.isSuccess shouldBe true + val (physical, stateTransfer) = result.get + stateTransfer shouldBe None + physical.opExecInitInfo match { + case OpExecWithCode(code, _) => code shouldBe "yield reconfigured" + case other => fail(s"expected OpExecWithCode, got $other") + } + } } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/r/RUDFOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/r/RUDFOpDescSpec.scala index 4963af8d6d6..fbffc5d6cc7 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/r/RUDFOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/udf/r/RUDFOpDescSpec.scala @@ -22,7 +22,8 @@ package org.apache.texera.amber.operator.udf.r import org.apache.texera.amber.core.executor.OpExecWithCode import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} -import org.apache.texera.amber.operator.LogicalOp +import org.apache.texera.amber.core.workflow.{PortIdentity, UnknownPartition} +import org.apache.texera.amber.operator.{LogicalOp, PortDescription} import org.apache.texera.amber.operator.metadata.OperatorGroupConstants import org.apache.texera.amber.util.JSONUtils.objectMapper import org.scalatest.flatspec.AnyFlatSpec @@ -109,4 +110,88 @@ class RUDFOpDescSpec extends AnyFlatSpec with Matchers { r.useTupleAPI shouldBe true r.retainInputColumns shouldBe true } + + "RUDFOpDesc.getPhysicalOp" should + "use r-tuple, parallelizable with the suggested worker count when workers > 1 and useTupleAPI" in { + val d = new RUDFOpDesc + d.code = "gen" + d.useTupleAPI = true + d.workers = 3 + val physical = d.getPhysicalOp(workflowId, executionId) + physical.opExecInitInfo match { + case OpExecWithCode(code, language) => + language shouldBe "r-tuple" + code shouldBe "gen" + case other => fail(s"expected OpExecWithCode, got $other") + } + physical.parallelizable shouldBe true + physical.suggestedWorkerNum shouldBe Some(3) + } + + it should "derive partitionRequirement from custom inputPorts when set" in { + val d = new RUDFOpDesc + d.code = "f" + d.inputPorts = List( + PortDescription( + portID = "0", + displayName = "in0", + disallowMultiInputs = false, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ) + ) + d.getPhysicalOp(workflowId, executionId).partitionRequirement shouldBe List( + Some(UnknownPartition()) + ) + } + + "RUDFOpDesc.operatorInfo" should + "build ports from custom inputPorts/outputPorts descriptions when provided" in { + val d = new RUDFOpDesc + d.inputPorts = List( + PortDescription( + portID = "0", + displayName = "in0", + disallowMultiInputs = true, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ) + ) + d.outputPorts = List( + PortDescription( + portID = "0", + displayName = "out0", + disallowMultiInputs = false, + isDynamicPort = false, + partitionRequirement = UnknownPartition(), + dependencies = List.empty + ) + ) + val info = d.operatorInfo + info.inputPorts should have length 1 + info.inputPorts.head.id shouldBe PortIdentity(0) + info.inputPorts.head.displayName shouldBe "in0" + info.inputPorts.head.disallowMultiLinks shouldBe true + info.outputPorts should have length 1 + info.outputPorts.head.displayName shouldBe "out0" + } + + "RUDFOpDesc.runtimeReconfiguration" should + "return the new op's physical op with no state transfer" in { + val oldOp = new RUDFOpDesc + val newOp = new RUDFOpDesc + newOp.code = "f2" + val result = newOp.runtimeReconfiguration(workflowId, executionId, oldOp, newOp) + result.isSuccess shouldBe true + val (physical, stateTransfer) = result.get + stateTransfer shouldBe None + physical.opExecInitInfo match { + case OpExecWithCode(code, language) => + code shouldBe "f2" + language shouldBe "r-table" + case other => fail(s"expected OpExecWithCode, got $other") + } + } }