Skip to content

Commit 9c8e335

Browse files
johnedquinnjohnedquinn
committed
Adds support for dynamic dispatch in eval
1 parent 1dabd4e commit 9c8e335

File tree

16 files changed

+247
-54
lines changed

16 files changed

+247
-54
lines changed

‎partiql-eval/build.gradle.kts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ dependencies {
3131
testImplementation(project(":partiql-parser"))
3232
testImplementation(project(":plugins:partiql-local"))
3333
testImplementation(project(":plugins:partiql-memory"))
34+
testImplementation(project(":plugins:partiql-plugin"))
3435
testImplementation(testFixtures(project(":partiql-planner")))
3536
testImplementation(testFixtures(project(":partiql-lang")))
3637
testImplementation(Deps.junit4)

‎partiql-eval/src/main/kotlin/org/partiql/eval/PartiQLEngineBuilder.kt

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,23 @@
11
package org.partiql.eval
22

33
import org.partiql.spi.connector.ConnectorBindings
4+
import org.partiql.spi.function.PartiQLFunction
5+
import org.partiql.spi.function.PartiQLFunctionExperimental
46

57
class PartiQLEngineBuilder {
68

79
private var catalogs: MutableMap<String, ConnectorBindings> = mutableMapOf()
810

11+
@OptIn(PartiQLFunctionExperimental::class)
12+
private var functions: MutableMap<String, List<PartiQLFunction>> = mutableMapOf()
13+
914
/**
1015
* Build the builder, return an implementation of a [PartiQLEngine]
1116
*
1217
* @return
1318
*/
14-
public fun build(): PartiQLEngine = PartiQLEngineDefault(catalogs)
19+
@OptIn(PartiQLFunctionExperimental::class)
20+
public fun build(): PartiQLEngine = PartiQLEngineDefault(catalogs, functions)
1521

1622
/**
1723
* Java style method for assigning a Catalog name to [ConnectorBindings].
@@ -33,4 +39,23 @@ class PartiQLEngineBuilder {
3339
public fun catalogs(vararg catalogs: Pair<String, ConnectorBindings>): PartiQLEngineBuilder = this.apply {
3440
this.catalogs = mutableMapOf(*catalogs)
3541
}
42+
43+
/**
44+
* Kotlin style method for assigning Catalog names to its implementations of [PartiQLFunction].
45+
*
46+
* @param catalogs
47+
* @return
48+
*/
49+
@OptIn(PartiQLFunctionExperimental::class)
50+
@Deprecated(
51+
message = """
52+
This will be replaced by PartiQL SPI's ability to gather functions at a specific path.
53+
However, this temporarily unblocks the injection of custom functions while function paths are worked on. This
54+
should only be used by the PartiQL library. External users: please consult with the PartiQL Maintainers
55+
if you intend on using this API.
56+
"""
57+
)
58+
public fun functions(vararg functions: Pair<String, List<PartiQLFunction>>): PartiQLEngineBuilder = this.apply {
59+
this.functions = mutableMapOf(*functions)
60+
}
3661
}

‎partiql-eval/src/main/kotlin/org/partiql/eval/PartiQLEngineDefault.kt

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,20 @@ import org.partiql.eval.internal.Compiler
44
import org.partiql.eval.internal.Record
55
import org.partiql.plan.PartiQLPlan
66
import org.partiql.spi.connector.ConnectorBindings
7+
import org.partiql.spi.function.PartiQLFunction
8+
import org.partiql.spi.function.PartiQLFunctionExperimental
79
import org.partiql.value.PartiQLValue
810
import org.partiql.value.PartiQLValueExperimental
911

10-
internal class PartiQLEngineDefault(
12+
internal class PartiQLEngineDefault @OptIn(PartiQLFunctionExperimental::class) constructor(
1113
private val catalogs: Map<String, ConnectorBindings>,
14+
private val functions: Map<String, List<PartiQLFunction>>
1215
) : PartiQLEngine {
1316

14-
@OptIn(PartiQLValueExperimental::class)
17+
@OptIn(PartiQLValueExperimental::class, PartiQLFunctionExperimental::class)
1518
override fun prepare(plan: PartiQLPlan): PartiQLStatement<*> {
1619
try {
17-
val compiler = Compiler(plan, catalogs)
20+
val compiler = Compiler(plan, catalogs, functions)
1821
val expression = compiler.compile()
1922
return object : PartiQLStatement.Query {
2023
override fun execute(): PartiQLValue {

‎partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ import org.partiql.eval.internal.operator.rel.RelJoinRight
1010
import org.partiql.eval.internal.operator.rel.RelProject
1111
import org.partiql.eval.internal.operator.rel.RelScan
1212
import org.partiql.eval.internal.operator.rel.RelScanIndexed
13+
import org.partiql.eval.internal.operator.rex.ExprCallDynamic
14+
import org.partiql.eval.internal.operator.rex.ExprCallStatic
1315
import org.partiql.eval.internal.operator.rex.ExprCase
1416
import org.partiql.eval.internal.operator.rex.ExprCollection
1517
import org.partiql.eval.internal.operator.rex.ExprGlobal
@@ -30,12 +32,16 @@ import org.partiql.plan.Statement
3032
import org.partiql.plan.visitor.PlanBaseVisitor
3133
import org.partiql.spi.connector.ConnectorBindings
3234
import org.partiql.spi.connector.ConnectorObjectPath
35+
import org.partiql.spi.function.PartiQLFunction
36+
import org.partiql.spi.function.PartiQLFunctionExperimental
37+
import org.partiql.types.function.FunctionSignature
3338
import org.partiql.value.PartiQLValueExperimental
3439
import java.lang.IllegalStateException
3540

36-
internal class Compiler(
41+
internal class Compiler @OptIn(PartiQLFunctionExperimental::class) constructor(
3742
private val plan: PartiQLPlan,
3843
private val catalogs: Map<String, ConnectorBindings>,
44+
private val functions: Map<String, List<PartiQLFunction>>
3945
) : PlanBaseVisitor<Operator, Unit>() {
4046

4147
fun compile(): Operator.Expr {
@@ -123,6 +129,33 @@ internal class Compiler(
123129
return ExprPathIndex(root, index)
124130
}
125131

132+
@OptIn(PartiQLFunctionExperimental::class)
133+
override fun visitRexOpCallStatic(node: Rex.Op.Call.Static, ctx: Unit): Operator {
134+
val function = getFunction(node.fn.signature)
135+
val args = node.args.map { visitRex(it, ctx) }.toTypedArray()
136+
return ExprCallStatic(function, args)
137+
}
138+
139+
@OptIn(PartiQLFunctionExperimental::class, PartiQLValueExperimental::class)
140+
override fun visitRexOpCallDynamic(node: Rex.Op.Call.Dynamic, ctx: Unit): Operator {
141+
val args = node.args.map { visitRex(it, ctx) }.toTypedArray()
142+
val candidates = node.candidates.map { candidate ->
143+
val fn = getFunction(candidate.fn.signature)
144+
val coercions = candidate.coercions.map { it?.signature?.let { sig -> getFunction(sig) } }
145+
ExprCallDynamic.Candidate(candidate.parameters.toTypedArray(), fn, coercions)
146+
}
147+
return ExprCallDynamic(candidates, args)
148+
}
149+
150+
@OptIn(PartiQLFunctionExperimental::class)
151+
private fun getFunction(signature: FunctionSignature): PartiQLFunction.Scalar {
152+
// TODO: .first() is a HACK. Once functions in the plan reference functions in a catalog, we will need to
153+
// query that connector. This should be a somewhat simple change.
154+
return functions.values.first().firstOrNull { function ->
155+
signature == function.signature
156+
} as PartiQLFunction.Scalar? ?: error("Could not find function signature: $signature")
157+
}
158+
126159
// REL
127160

128161
override fun visitRel(node: Rel, ctx: Unit): Operator.Relation {

‎partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCallDynamic.kt

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package org.partiql.eval.internal.operator.rex
2+
3+
import org.partiql.eval.internal.Record
4+
import org.partiql.eval.internal.operator.Operator
5+
import org.partiql.spi.function.PartiQLFunction
6+
import org.partiql.spi.function.PartiQLFunctionExperimental
7+
import org.partiql.value.PartiQLValue
8+
import org.partiql.value.PartiQLValueExperimental
9+
import org.partiql.value.PartiQLValueType
10+
import org.partiql.value.missingValue
11+
12+
/**
13+
* This represents Dynamic Dispatch.
14+
*
15+
* For the purposes of efficiency, this implementation aims to reduce any re-execution of compiled arguments. It
16+
* does this by avoiding the compilation of [Candidate.fn] and [Candidate.coercions] into
17+
* [ExprCallStatic]'s. By doing this, this implementation can evaluate ([eval]) the input [Record], execute and gather the
18+
* arguments, and pass the [PartiQLValue]s directly to the [Candidate.eval].
19+
*/
20+
internal class ExprCallDynamic(
21+
private val candidates: List<Candidate>,
22+
private val args: Array<Operator.Expr>
23+
) : Operator.Expr {
24+
25+
@OptIn(PartiQLValueExperimental::class)
26+
override fun eval(record: Record): PartiQLValue {
27+
val actualArgs = args.map { it.eval(record) }.toTypedArray()
28+
candidates.forEach { candidate ->
29+
if (candidate.matches(actualArgs)) {
30+
candidate.eval(actualArgs)
31+
}
32+
}
33+
return missingValue()
34+
}
35+
36+
/**
37+
* This represents a single candidate for dynamic dispatch.
38+
*
39+
* This implementation assumes that the [eval] input [Record] contains the original arguments for the desired [fn].
40+
* It performs the coercions (if necessary) before computing the result.
41+
*
42+
* @see ExprCallDynamic
43+
*/
44+
internal class Candidate @OptIn(PartiQLValueExperimental::class, PartiQLFunctionExperimental::class) constructor(
45+
val types: Array<PartiQLValueType>,
46+
val fn: PartiQLFunction.Scalar,
47+
val coercions: List<PartiQLFunction.Scalar?>
48+
) {
49+
50+
@OptIn(PartiQLValueExperimental::class, PartiQLFunctionExperimental::class)
51+
fun eval(originalArgs: Array<PartiQLValue>): PartiQLValue {
52+
val args = coercions.mapIndexed { index, coercion ->
53+
coercion?.invoke(arrayOf(originalArgs[index])) ?: originalArgs[index]
54+
}.toTypedArray()
55+
return fn.invoke(args)
56+
}
57+
58+
@OptIn(PartiQLValueExperimental::class)
59+
internal fun matches(args: Array<PartiQLValue>): Boolean {
60+
for (i in args.indices) {
61+
if (args[i].type != types[i]) {
62+
return false
63+
}
64+
}
65+
return true
66+
}
67+
}
68+
}

‎partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCall.kt renamed to ‎partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprCallStatic.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import org.partiql.value.PartiQLValueExperimental
1111
import org.partiql.value.missingValue
1212

1313
@OptIn(PartiQLValueExperimental::class, PartiQLFunctionExperimental::class)
14-
internal class ExprCall(
14+
internal class ExprCallStatic(
1515
private val fn: PartiQLFunction.Scalar,
1616
private val inputs: Array<Operator.Expr>,
1717
) : Operator.Expr {

‎partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt

Lines changed: 47 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import org.partiql.eval.PartiQLResult
1111
import org.partiql.parser.PartiQLParser
1212
import org.partiql.planner.PartiQLPlanner
1313
import org.partiql.planner.PartiQLPlannerBuilder
14+
import org.partiql.plugin.PartiQLPlugin
15+
import org.partiql.spi.function.PartiQLFunctionExperimental
1416
import org.partiql.value.PartiQLValue
1517
import org.partiql.value.PartiQLValueExperimental
1618
import org.partiql.value.bagValue
@@ -212,6 +214,18 @@ class PartiQLEngineDefaultTest {
212214
input = "SELECT DISTINCT VALUE t FROM <<true, false, true, false, false, false>> AS t;",
213215
expected = bagValue(boolValue(true), boolValue(false))
214216
),
217+
SuccessTestCase(
218+
input = "SELECT DISTINCT VALUE t FROM <<true, false, true, false, false, false>> AS t WHERE t = TRUE;",
219+
expected = bagValue(boolValue(true))
220+
),
221+
SuccessTestCase(
222+
input = "100 + 50;",
223+
expected = int32Value(150)
224+
),
225+
SuccessTestCase(
226+
input = "SELECT DISTINCT VALUE t * 100 FROM <<0, 1, 2, 3>> AS t;",
227+
expected = bagValue(int32Value(0), int32Value(100), int32Value(200), int32Value(300))
228+
),
215229
SuccessTestCase(
216230
input = """
217231
PIVOT x.v AT x.k FROM <<
@@ -226,14 +240,38 @@ class PartiQLEngineDefaultTest {
226240
"c" to stringValue("z"),
227241
)
228242
),
243+
SuccessTestCase(
244+
input = """
245+
CASE (1)
246+
WHEN NULL THEN 'isNull'
247+
WHEN MISSING THEN 'isMissing'
248+
WHEN 2 THEN 'isTwo'
249+
END
250+
;
251+
""".trimIndent(),
252+
expected = nullValue()
253+
),
254+
SuccessTestCase(
255+
input = """
256+
CASE (1)
257+
WHEN NULL THEN 'isNull'
258+
WHEN MISSING THEN 'isMissing'
259+
WHEN 2 THEN 'isTwo'
260+
WHEN 1 THEN 'isOne'
261+
END
262+
;
263+
""".trimIndent(),
264+
expected = stringValue("isOne")
265+
)
229266
)
230267
}
231268
public class SuccessTestCase @OptIn(PartiQLValueExperimental::class) constructor(
232269
val input: String,
233270
val expected: PartiQLValue
234271
) {
235272

236-
private val engine = PartiQLEngine.default()
273+
@OptIn(PartiQLFunctionExperimental::class)
274+
private val engine = PartiQLEngine.builder().functions("partiql" to PartiQLPlugin.functions).build()
237275
private val planner = PartiQLPlannerBuilder().build()
238276
private val parser = PartiQLParser.default()
239277

@@ -266,32 +304,17 @@ class PartiQLEngineDefaultTest {
266304
}
267305
}
268306

269-
@Disabled("This is disabled because FN EQUALS is not yet implemented.")
270307
@Test
271-
fun testCaseLiteral02() = SuccessTestCase(
272-
input = """
273-
CASE (1)
274-
WHEN NULL THEN 'isNull'
275-
WHEN MISSING THEN 'isMissing'
276-
WHEN 2 THEN 'isTwo'
277-
WHEN 1 THEN 'isOne'
278-
END
279-
;
280-
""".trimIndent(),
281-
expected = stringValue("isOne")
308+
@Disabled("CASTS have not yet been implemented.")
309+
fun testCast1() = SuccessTestCase(
310+
input = "1 + 2.0",
311+
expected = int32Value(3),
282312
).assert()
283313

284-
@Disabled("This is disabled because FN EQUALS is not yet implemented.")
285314
@Test
286-
fun testCaseLiteral03() = SuccessTestCase(
287-
input = """
288-
CASE (1)
289-
WHEN NULL THEN 'isNull'
290-
WHEN MISSING THEN 'isMissing'
291-
WHEN 2 THEN 'isTwo'
292-
END
293-
;
294-
""".trimIndent(),
295-
expected = nullValue()
315+
@Disabled("CASTS have not yet been implemented.")
316+
fun testCasts() = SuccessTestCase(
317+
input = "SELECT DISTINCT VALUE t * 100 FROM <<0, 1, 2.0, 3.0>> AS t;",
318+
expected = bagValue(int32Value(0), int32Value(100), int32Value(200), int32Value(300))
296319
).assert()
297320
}

‎partiql-plan/src/main/resources/partiql_plan.ion

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
imports::{
22
kotlin: [
33
partiql_value::'org.partiql.value.PartiQLValue',
4+
partiql_value_type::'org.partiql.value.PartiQLValueType',
45
static_type::'org.partiql.types.StaticType',
56
scalar_signature::'org.partiql.types.function.FunctionSignature$Scalar',
67
aggregation_signature::'org.partiql.types.function.FunctionSignature$Aggregation',
@@ -116,8 +117,16 @@ rex::{
116117
args: list::[rex],
117118
candidates: list::[candidate],
118119
_: [
120+
// Represents a potential candidate for dynamic dispatch. AKA `SELECT abs(a) FROM << 1, 2.0 >> AS a` can invoke
121+
// ABS(INT32) -> INT32 or ABS(DEC) -> DEC. In this scenario, we maintain the two potential candidates.
122+
//
123+
// @param fn - represents the function to invoke (ex: ABS(INT32) -> INT32)
124+
// @param parameters - represents the input type(s) to match. (ex: INT32)
125+
// @param coercions - represents the optional coercion to use on the argument(s). It will be NULL if no coercion
126+
// is necessary.
119127
candidate::{
120128
fn: fn,
129+
parameters: list::[partiql_value_type],
121130
coercions: list::[optional::fn]
122131
}
123132
]

‎partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Nodes.kt

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ import org.partiql.types.StaticType
6565
import org.partiql.types.function.FunctionSignature
6666
import org.partiql.value.PartiQLValue
6767
import org.partiql.value.PartiQLValueExperimental
68+
import org.partiql.value.PartiQLValueType
6869
import kotlin.random.Random
6970

7071
internal abstract class PlanNode {
@@ -550,9 +551,11 @@ internal data class Rex(
550551

551552
internal data class Candidate(
552553
@JvmField
553-
internal val fn: Fn.Resolved,
554+
internal val fn: Fn,
554555
@JvmField
555-
internal val coercions: List<Fn.Resolved?>,
556+
internal val parameters: List<PartiQLValueType>,
557+
@JvmField
558+
internal val coercions: List<Fn?>,
556559
) : PlanNode() {
557560
internal override val children: List<PlanNode> by lazy {
558561
val kids = mutableListOf<PlanNode?>()

‎partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Plan.kt

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import org.partiql.types.StaticType
77
import org.partiql.types.function.FunctionSignature
88
import org.partiql.value.PartiQLValue
99
import org.partiql.value.PartiQLValueExperimental
10+
import org.partiql.value.PartiQLValueType
1011

1112
internal fun partiQLPlan(
1213
version: PartiQLVersion,
@@ -69,8 +70,11 @@ internal fun rexOpCallStatic(fn: Fn, args: List<Rex>): Rex.Op.Call.Static = Rex.
6970
internal fun rexOpCallDynamic(args: List<Rex>, candidates: List<Rex.Op.Call.Dynamic.Candidate>):
7071
Rex.Op.Call.Dynamic = Rex.Op.Call.Dynamic(args, candidates)
7172

72-
internal fun rexOpCallDynamicCandidate(fn: Fn.Resolved, coercions: List<Fn.Resolved?>):
73-
Rex.Op.Call.Dynamic.Candidate = Rex.Op.Call.Dynamic.Candidate(fn, coercions)
73+
internal fun rexOpCallDynamicCandidate(
74+
fn: Fn,
75+
parameters: List<PartiQLValueType>,
76+
coercions: List<Fn?>,
77+
): Rex.Op.Call.Dynamic.Candidate = Rex.Op.Call.Dynamic.Candidate(fn, parameters, coercions)
7478

7579
internal fun rexOpCase(branches: List<Rex.Op.Case.Branch>, default: Rex): Rex.Op.Case =
7680
Rex.Op.Case(branches, default)

0 commit comments

Comments
 (0)