Skip to content

Commit

Permalink
Merge 5d676bb into f2369ba
Browse files Browse the repository at this point in the history
  • Loading branch information
yliuuuu authored Feb 29, 2024
2 parents f2369ba + 5d676bb commit 3d2aa2e
Show file tree
Hide file tree
Showing 5 changed files with 271 additions and 42 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1289,4 +1289,109 @@ class PartiQLEngineDefaultTest {
""".trimIndent(),
expected = boolValue(false)
).assert()

@Test
// TODO: Add to conformance tests
fun wildCard() =
SuccessTestCase(
input = """
[
{ 'id':'5',
'books':[
{ 'title':'A',
'price':5.0,
'authors': [{'name': 'John'}, {'name': 'Doe'}]
},
{ 'title':'B',
'price':2.0,
'authors': [{'name': 'Zoe'}, {'name': 'Bill'}]
}
]
},
{ 'id':'6',
'books':[
{ 'title':'A',
'price':5.0,
'authors': [{'name': 'John'}, {'name': 'Doe'}]
},
{ 'title':'E',
'price':2.0,
'authors': [{'name': 'Zoe'}, {'name': 'Bill'}]
}
]
},
{ 'id':7,
'books':[]
}
][*].books[*].authors[*].name
""".trimIndent(),
expected = bagValue(
listOf(
stringValue("John"), stringValue("Doe"), stringValue("Zoe"), stringValue("Bill"),
stringValue("John"), stringValue("Doe"), stringValue("Zoe"), stringValue("Bill")
)
)
).assert()

@Test
// TODO: add to conformance tests
// Note that the existing pipeline produced identical result when supplying with
// SELECT VALUE v2.name FROM e as v0, v0.books as v1, unpivot v1.authors as v2;
// But it produces different result when supplying with e[*].books[*].authors.*
// <<
// <<{ 'name': 'John'},{'name': 'Doe'} >>,
// ...
// >>
fun unpivot() =
SuccessTestCase(
input = """
[
{ 'id':'5',
'books':[
{ 'title':'A',
'price':5.0,
'authors': {
'first': {'name': 'John'},
'second': {'name': 'Doe'}
}
},
{ 'title':'B',
'price':2.0,
'authors': {
'first': {'name': 'Zoe'},
'second': {'name': 'Bill'}
}
}
]
},
{ 'id':'6',
'books':[
{ 'title':'A',
'price':5.0,
'authors': {
'first': {'name': 'John'},
'second': {'name': 'Doe'}
}
},
{ 'title':'E',
'price':2.0,
'authors': {
'first': {'name': 'Zoe'},
'second': {'name': 'Bill'}
}
}
]
},
{ 'id':7,
'books':[]
}
][*].books[*].authors.*.name
""".trimIndent(),
expected = bagValue(
listOf(
stringValue("John"), stringValue("Doe"), stringValue("Zoe"), stringValue("Bill"),
stringValue("John"), stringValue("Doe"), stringValue("Zoe"), stringValue("Bill")
)
)
).assert()
}
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,6 @@ internal object RelConverter {
}

// Helpers

private fun convertScan(rex: Rex, binding: Rel.Binding): Rel {
val schema = listOf(binding)
val props = emptySet<Rel.Prop>()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,17 @@ import org.partiql.ast.Type
import org.partiql.ast.visitor.AstBaseVisitor
import org.partiql.planner.internal.Env
import org.partiql.planner.internal.ir.Identifier
import org.partiql.planner.internal.ir.Rel
import org.partiql.planner.internal.ir.Rex
import org.partiql.planner.internal.ir.builder.plan
import org.partiql.planner.internal.ir.identifierQualified
import org.partiql.planner.internal.ir.identifierSymbol
import org.partiql.planner.internal.ir.rel
import org.partiql.planner.internal.ir.relBinding
import org.partiql.planner.internal.ir.relOpJoin
import org.partiql.planner.internal.ir.relOpScan
import org.partiql.planner.internal.ir.relOpUnpivot
import org.partiql.planner.internal.ir.relType
import org.partiql.planner.internal.ir.rex
import org.partiql.planner.internal.ir.rexOpCallUnresolved
import org.partiql.planner.internal.ir.rexOpCastUnresolved
Expand All @@ -36,17 +43,20 @@ import org.partiql.planner.internal.ir.rexOpLit
import org.partiql.planner.internal.ir.rexOpPathIndex
import org.partiql.planner.internal.ir.rexOpPathKey
import org.partiql.planner.internal.ir.rexOpPathSymbol
import org.partiql.planner.internal.ir.rexOpSelect
import org.partiql.planner.internal.ir.rexOpStruct
import org.partiql.planner.internal.ir.rexOpStructField
import org.partiql.planner.internal.ir.rexOpSubquery
import org.partiql.planner.internal.ir.rexOpTupleUnion
import org.partiql.planner.internal.ir.rexOpVarLocal
import org.partiql.planner.internal.ir.rexOpVarUnresolved
import org.partiql.planner.internal.typer.toNonNullStaticType
import org.partiql.planner.internal.typer.toStaticType
import org.partiql.types.StaticType
import org.partiql.value.PartiQLValueExperimental
import org.partiql.value.PartiQLValueType
import org.partiql.value.StringValue
import org.partiql.value.boolValue
import org.partiql.value.int32Value
import org.partiql.value.int64Value
import org.partiql.value.io.PartiQLValueIonReaderBuilder
Expand Down Expand Up @@ -249,41 +259,152 @@ internal object RexConverter {
else -> root to node.steps
}

// Return wrapped path
return when (newSteps.isEmpty()) {
true -> newRoot
false -> newSteps.fold(newRoot) { current, step ->
val path = when (step) {
is Expr.Path.Step.Index -> {
val key = visitExprCoerce(step.key, context)
when (val astKey = step.key) {
is Expr.Lit -> when (astKey.value) {
is StringValue -> rexOpPathKey(current, key)
else -> rexOpPathIndex(current, key)
}
is Expr.Cast -> when (astKey.asType is Type.String) {
true -> rexOpPathKey(current, key)
false -> rexOpPathIndex(current, key)
}
if (newSteps.isEmpty()) {
return newRoot
}

val fromList = mutableListOf<Rel>()

var varRefIndex = 0 // tracking var ref index

val pathNavi = newSteps.fold(newRoot) { current, step ->
val path = when (step) {
is Expr.Path.Step.Index -> {
val key = visitExprCoerce(step.key, context)
val op = when (val astKey = step.key) {
is Expr.Lit -> when (astKey.value) {
is StringValue -> rexOpPathKey(current, key)
else -> rexOpPathIndex(current, key)
}
}
is Expr.Path.Step.Symbol -> {
val identifier = AstToPlan.convert(step.symbol)
when (identifier.caseSensitivity) {
Identifier.CaseSensitivity.SENSITIVE -> rexOpPathKey(
current,
rexString(identifier.symbol)
)
Identifier.CaseSensitivity.INSENSITIVE -> rexOpPathSymbol(current, identifier.symbol)

is Expr.Cast -> when (astKey.asType is Type.String) {
true -> rexOpPathKey(current, key)
false -> rexOpPathIndex(current, key)
}

else -> rexOpPathIndex(current, key)
}
is Expr.Path.Step.Unpivot -> error("Unpivot path not supported yet")
is Expr.Path.Step.Wildcard -> error("Wildcard path not supported yet")
op
}

is Expr.Path.Step.Symbol -> {
val identifier = AstToPlan.convert(step.symbol)
val op = when (identifier.caseSensitivity) {
Identifier.CaseSensitivity.SENSITIVE -> rexOpPathKey(
current,
rexString(identifier.symbol)
)

Identifier.CaseSensitivity.INSENSITIVE -> rexOpPathSymbol(current, identifier.symbol)
}
op
}

// Unpivot and Wildcard steps trigger the rewrite
// According to spec Section 4.3
// ew1p1...wnpn
// rewrite to:
// SELECT VALUE v_n.p_n
// FROM
// u_1 e as v_1
// u_2 @v_1.p_1 as v_2
// ...
// u_n @v_(n-1).p_(n-1) as v_n
// The From clause needs to be rewritten to
// Join <------------------- schema: [(k_1), v_1, (k_2), v_2, ..., (k_(n-1)) v_(n-1)]
// / \
// ... un @v_(n-1).p_(n-1) <-- stack: [global, typeEnv: [outer: [global], schema: [(k_1), v_1, (k_2), v_2, ..., (k_(n-1)) v_(n-1)]]]
// Join <----------------------- schema: [(k_1), v_1, (k_2), v_2, (k_3), v_3]
// / \
// u_2 @v_1.p_1 as v2 <------- stack: [global, typeEnv: [outer: [global], schema: [(k_1), v_1, (k_2), v_2]]]
// JOIN <---------------------------- schema: [(k_1), v_1, (k_2), v_2]
// / \
// u_1 e as v_1 < ----\----------------------- stack: [global]
// u_2 @v_1.p_1 as v2 <------ stack: [global, typeEnv: [outer: [global], schema: [(k_1), v_1]]]
// while doing the traversal, instead of passing the stack,
// each join will produce its own schema and pass the schema as a type Env.
// The (k_i) indicate the possible key binding produced by unpivot.
// We calculate the var ref on the fly.
is Expr.Path.Step.Unpivot -> {
// Unpivot produces two binding, in this context we want the value,
// which always going to be the second binding
val op = rexOpVarLocal(1, varRefIndex + 1)
varRefIndex += 2
val index = fromList.size
fromList.add(relFromUnpivot(current, index))
op
}
is Expr.Path.Step.Wildcard -> {
// Scan produce only one binding
val op = rexOpVarLocal(1, varRefIndex)
varRefIndex += 1
val index = fromList.size
fromList.add(relFromDefault(current, index))
op
}
rex(StaticType.ANY, path)
}
rex(StaticType.ANY, path)
}

if (fromList.size == 0) return pathNavi
val fromNode = fromList.reduce { acc, scan ->
val schema = acc.type.schema + scan.type.schema
val props = emptySet<Rel.Prop>()
val type = relType(schema, props)
rel(type, relOpJoin(acc, scan, rex(StaticType.BOOL, rexOpLit(boolValue(true))), Rel.Op.Join.Type.INNER))
}

// compute the ref used by select construct
// always going to be the last binding
val selectRef = fromNode.type.schema.size - 1

val constructor = when (val op = pathNavi.op) {
is Rex.Op.Path.Index -> rex(pathNavi.type, rexOpPathIndex(rex(op.root.type, rexOpVarLocal(0, selectRef)), op.key))
is Rex.Op.Path.Key -> rex(pathNavi.type, rexOpPathKey(rex(op.root.type, rexOpVarLocal(0, selectRef)), op.key))
is Rex.Op.Path.Symbol -> rex(pathNavi.type, rexOpPathSymbol(rex(op.root.type, rexOpVarLocal(0, selectRef)), op.key))
is Rex.Op.Var.Local -> rex(pathNavi.type, rexOpVarLocal(0, selectRef))
else -> throw IllegalStateException()
}
val op = rexOpSelect(constructor, fromNode)
return rex(StaticType.ANY, op)
}

/**
* Construct Rel(Scan([path])).
*
* The constructed rel would produce one binding: _v$[index]
*/
private fun relFromDefault(path: Rex, index: Int): Rel {
val schema = listOf(
relBinding(
name = "_v$index", // fresh variable
type = path.type
)
)
val props = emptySet<Rel.Prop>()
val relType = relType(schema, props)
return rel(relType, relOpScan(path))
}

/**
* Construct Rel(Unpivot([path])).
*
* The constructed rel would produce two bindings: _k$[index] and _v$[index]
*/
private fun relFromUnpivot(path: Rex, index: Int): Rel {
val schema = listOf(
relBinding(
name = "_k$index", // fresh variable
type = StaticType.STRING
),
relBinding(
name = "_v$index", // fresh variable
type = path.type
)
)
val props = emptySet<Rel.Prop>()
val relType = relType(schema, props)
return rel(relType, relOpUnpivot(path))
}

private fun rexString(str: String) = rex(StaticType.STRING, rexOpLit(stringValue(str)))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -145,11 +145,6 @@ internal class PlanTyper(
return rel(type, op)
}

override fun visitRelOpErr(node: Rel.Op.Err, ctx: Rel.Type?): Rel {
val type = ctx ?: relType(emptyList(), emptySet())
return rel(type, node)
}

/**
* The output schema of a `rel.op.scan_index` is the value binding and index binding.
*/
Expand All @@ -176,15 +171,19 @@ internal class PlanTyper(
val kType = STRING

// value type, possibly coerced.
val vType = when (val t = rex.type) {
is StructType -> {
if (t.contentClosed || t.constraints.contains(TupleConstraint.Open(false))) {
unionOf(t.fields.map { it.value }.toSet()).flatten()
} else {
ANY
val vType = rex.type.allTypes.map { type ->
when (type) {
is StructType -> {
if (type.contentClosed || type.constraints.contains(TupleConstraint.Open(false))) {
unionOf(type.fields.map { it.value }.toSet()).flatten()
} else {
ANY
}
}
else -> type
}
else -> t
}.let {
unionOf(it.toSet()).flatten()
}

// rewrite
Expand All @@ -193,6 +192,11 @@ internal class PlanTyper(
return rel(type, op)
}

override fun visitRelOpErr(node: Rel.Op.Err, ctx: Rel.Type?): Rel {
val type = ctx ?: relType(emptyList(), emptySet())
return rel(type, node)
}

override fun visitRelOpDistinct(node: Rel.Op.Distinct, ctx: Rel.Type?): Rel {
val input = visitRel(node.input, ctx)
return rel(input.type, relOpDistinct(input))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ internal data class TypeEnv(
internal fun getScope(depth: Int): TypeEnv {
return when (depth) {
0 -> this
else -> outer.reversed()[depth - 1]
else -> outer[outer.size - depth]
}
}

Expand Down

0 comments on commit 3d2aa2e

Please sign in to comment.