Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@yliuuuu
yliuuuu committed Dec 25, 2024
1 parent 3f7a1e4 commit 14617d0
Show file tree
Hide file tree
Showing 19 changed files with 680 additions and 151 deletions.
2 changes: 1 addition & 1 deletion partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Nodes.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -347,7 +347,7 @@ internal sealed class Statement : PlanNode() {
@JvmField
public val name: String,
@JvmField
public val `value`: PartiQLValue,
public val `value`: String,
) : PlanNode() {
public override val children: List<PlanNode> = emptyList()

Expand Down
4 changes: 1 addition & 3 deletions ...iql-planner/src/main/kotlin/org/partiql/planner/internal/normalize/NormalizeFromSource.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@ import org.partiql.ast.FromType
import org.partiql.ast.Query
import org.partiql.ast.QueryBody
import org.partiql.ast.Statement
import org.partiql.ast.ddl.Ddl
import org.partiql.ast.expr.Expr
import org.partiql.planner.internal.helpers.toBinder

Expand All @@ -35,12 +34,11 @@ import org.partiql.planner.internal.helpers.toBinder
*/
internal object NormalizeFromSource : AstPass {

override fun apply(statement: Statement): Statement = when(statement) {
override fun apply(statement: Statement): Statement = when (statement) {
is Query -> statement.accept(Visitor, 0) as Statement
else -> statement
}


private object Visitor : AstRewriter<Int>() {

// Each SFW starts the ctx count again.
Expand Down
2 changes: 1 addition & 1 deletion partiql-planner/src/main/kotlin/org/partiql/planner/internal/normalize/NormalizeGroupBy.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ import org.partiql.planner.internal.helpers.toBinder
*/
internal object NormalizeGroupBy : AstPass {

override fun apply(statement: Statement) = when(statement) {
override fun apply(statement: Statement) = when (statement) {
is Query -> Visitor.visitStatement(statement, 0) as Statement
else -> statement
}
Expand Down
1 change: 0 additions & 1 deletion partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/AstToPlan.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import org.partiql.ast.AstVisitor
import org.partiql.ast.DataType
import org.partiql.ast.Query
import org.partiql.ast.ddl.CreateTable
import org.partiql.ast.ddl.Ddl
import org.partiql.ast.expr.ExprQuerySet
import org.partiql.errors.TypeCheckException
import org.partiql.planner.internal.Env
Expand Down
25 changes: 20 additions & 5 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/DdlConverter.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ import org.partiql.ast.ddl.AttributeConstraint.Unique
import org.partiql.ast.ddl.ColumnDefinition
import org.partiql.ast.ddl.CreateTable
import org.partiql.ast.ddl.Ddl
import org.partiql.ast.ddl.KeyValue
import org.partiql.ast.ddl.PartitionBy
import org.partiql.ast.ddl.TableConstraint
import org.partiql.ast.sql.sql
import org.partiql.planner.internal.DdlField
Expand All @@ -20,6 +22,8 @@ import org.partiql.planner.internal.ir.statementDDL
import org.partiql.planner.internal.ir.statementDDLAttribute
import org.partiql.planner.internal.ir.statementDDLCommandCreateTable
import org.partiql.planner.internal.ir.statementDDLConstraintCheck
import org.partiql.planner.internal.ir.statementDDLPartitionByAttrList
import org.partiql.planner.internal.ir.statementDDLTableProperty
import org.partiql.planner.internal.transforms.AstToPlan.convert
import org.partiql.planner.internal.transforms.AstToPlan.visitType
import org.partiql.planner.internal.typer.CompilerType
Expand Down Expand Up @@ -75,10 +79,10 @@ internal object DdlConverter {
val pk = node.constraints.filterIsInstance<TableConstraint.Unique>()
.filter { it.isPrimaryKey }
.let {
when(it.size) {
when (it.size) {
0 -> emptyList()
1 -> it.first().columns.map { convert(it) }
else -> throw IllegalArgumentException("multiple PK")
else -> throw IllegalArgumentException("multiple PK")
}
}

Expand All @@ -88,18 +92,21 @@ internal object DdlConverter {
acc + constr.columns.map { convert(it) }
}

val partitionBy = node.partitionBy?.let { visitPartitionBy(it, ctx) }
val tableProperty = node.tableProperties.map { visitKeyValue(it, ctx) }

return statementDDLCommandCreateTable(
tableName,
attributes,
emptyList(),
null,
emptyList(),
partitionBy,
tableProperty,
pk,
unique
)
}

// !!! The planning stage ignores constraint name for now
// !!! The planning stage ignores the constraint name for now
override fun visitColumnDefinition(node: ColumnDefinition, ctx: Env): Statement.DDL.Attribute {
val name = convert(node.name)
val type = visitType(node.dataType, ctx)
Expand Down Expand Up @@ -156,6 +163,14 @@ internal object DdlConverter {
node.searchCondition.sql()
)

override fun visitPartitionBy(node: PartitionBy, ctx: Env): Statement.DDL.PartitionBy {
return statementDDLPartitionByAttrList(node.columns.map { convert(it) })
}

override fun visitKeyValue(node: KeyValue, ctx: Env): Statement.DDL.TableProperty {
return statementDDLTableProperty(node.key, node.value)
}

private fun visitType(node: DataType, ctx: Env): CompilerType {
// Struct requires special process in DDL
return if (node.code() == DataType.STRUCT) {
Expand Down
6 changes: 4 additions & 2 deletions ...src/main/kotlin/org/partiql/planner/internal/transforms/InlineCheckConstraintExtractor.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,15 +28,17 @@ internal object InlineCheckConstraintExtractor : PlanBaseVisitor<PShape, PShape>
override fun visitStatementDDLConstraintCheck(node: Statement.DDL.Constraint.Check, ctx: PShape): PShape {
val lowered = visitRex(node.expression, ctx)
// No lowering happened, then wrap the PShape with a generic Constraint trait
return if (lowered.equals(ctx)) { ConstraintTrait(ctx, node.sql) } else lowered
return if (lowered == ctx) {
ConstraintTrait(ctx, node.sql)
} else lowered
}

override fun visitRex(node: Rex, ctx: PShape): PShape {
return node.op.accept(this, ctx)
}

override fun visitRexOpCallStatic(node: Rex.Op.Call.Static, ctx: PShape): PShape {
return when(node.fn.name) {
return when (node.fn.name) {
"gte" -> getRhsAsNumericOrNull(node.args[1])?.let {
RangeTrait(ctx, it, null)
} ?: ctx
Expand Down
51 changes: 36 additions & 15 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/PlanTransform.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ internal class PlanTransform(private val flags: Set<PlannerFlag>) {
*/
fun transform(internal: IPlan, listener: PErrorListener): Plan {
val signal = flags.contains(PlannerFlag.SIGNAL_MODE)
when(internal.statement) {
when (internal.statement) {
is Statement.DDL -> {
val query = internal.statement
val visitor = DDLVisitor(listener, signal)
Expand Down Expand Up @@ -94,19 +94,42 @@ internal class PlanTransform(private val flags: Set<PlannerFlag>) {
override fun visitStatementDDLCommandCreateTable(node: Statement.DDL.Command.CreateTable, ctx: Unit): Action.CreateTable {
val fields = node.attributes.map {
val shape = visitStatementDDLAttribute(it, ctx)
Field.of(it.name.getIdentifier().getText(), shape)
if (node.primaryKey.contains(it.name)) {
Field.of(it.name.getIdentifier().getText(), NotNullTrait(shape))
} else {
Field.of(it.name.getIdentifier().getText(), shape)
}
}
val row = PShape(PType.row(fields))
.let {
if (node.primaryKey.isNotEmpty())
PrimaryKeyTrait(it, node.primaryKey.map { it.getIdentifier().getText() })
else it
}.let {
if (node.unique.isNotEmpty())
UniqueTrait(it, node.unique.map { it.getIdentifier().getText() })
else it
val schema = PShape(PType.bag(row)).let {
if (node.primaryKey.isNotEmpty())
PrimaryKeyTrait(it, node.primaryKey.map { it.getIdentifier().getText() })
else it
}.let {
if (node.unique.isNotEmpty())
UniqueTrait(it, node.unique.map { it.getIdentifier().getText() })
else it
}.let { it ->
var shape = it
if (node.tableProperties.isNotEmpty()) {
node.tableProperties.forEach { prop ->
shape = MetadataTrait(shape, prop.name, prop.value)
}
}
shape
}.let { it ->
when (val partition = node.partitionBy) {
is Statement.DDL.PartitionBy.AttrList -> {
val names = buildString {
append("[")
append(partition.attrs.joinToString(",") { it.getIdentifier().getText() })
append("]")
}
MetadataTrait(it, "partition", names)
}
null -> it
}
val schema = PShape(PType.bag(row))
}
return Action.CreateTable {
Table.builder()
.name(node.name.getIdentifier().getText())
Expand All @@ -115,12 +138,12 @@ internal class PlanTransform(private val flags: Set<PlannerFlag>) {
}
}

override fun visitStatementDDLAttribute(node: Statement.DDL.Attribute, ctx: Unit): PType {
override fun visitStatementDDLAttribute(node: Statement.DDL.Attribute, ctx: Unit): PShape {
val ddlField = DdlField.fromAttr(node)
return visitDdlField(ddlField, ctx)
}

private fun visitDdlField(ddlField: DdlField, ctx: Unit): PType {
private fun visitDdlField(ddlField: DdlField, ctx: Unit): PShape {
val baseShape = ddlField.type.getDelegate()
val fieldReduced = when (baseShape.code()) {
PType.ROW -> {
Expand All @@ -143,10 +166,8 @@ internal class PlanTransform(private val flags: Set<PlannerFlag>) {
} ?: constraintReduced
return metadataReduced
}

}


private class Visitor(
private val listener: PErrorListener,
private val signal: Boolean,
Expand Down
22 changes: 17 additions & 5 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/PlanTyper.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ import org.partiql.planner.internal.ir.rexOpStructField
import org.partiql.planner.internal.ir.rexOpSubquery
import org.partiql.planner.internal.ir.statementDDL
import org.partiql.planner.internal.ir.statementDDLCommandCreateTable
import org.partiql.planner.internal.ir.statementDDLPartitionByAttrList
import org.partiql.planner.internal.ir.statementQuery
import org.partiql.planner.internal.ir.util.PlanRewriter
import org.partiql.spi.Context
Expand Down Expand Up @@ -1285,22 +1286,32 @@ internal class PlanTyper(private val env: Env, config: Context) {
isDeclaredAttribute(uniqueAttr, attrs.map { DdlField.fromAttr(it) }) ?: throw IllegalArgumentException("Unresolved ref")
}

val tablePrimaryContr = node.primaryKey.map { uniqueAttr ->
isDeclaredAttribute(uniqueAttr, attrs.map { DdlField.fromAttr(it) }) ?: throw IllegalArgumentException("Unresolved ref")
val tablePrimaryContr = node.primaryKey.map { pkAtrr ->
isDeclaredAttribute(pkAtrr, attrs.map { DdlField.fromAttr(it) }) ?: throw IllegalArgumentException("Unresolved ref")
}

// ALSO: For PK
// Thing like PRIMARY KEY (FOO, FOO) will be rejected
if (tablePrimaryContr.toSet().size != tableUniqueContr.size) {
if (tablePrimaryContr.toSet().size != tablePrimaryContr.size) {
throw IllegalArgumentException("Attribute appears multiple times in primary key constraint")
}
val finalPK = tablePrimaryContr + attributePK
val finalUnique = (tableUniqueContr + attributeUnique).toSet().toList()
val finalUnique = (tableUniqueContr + attributeUnique + finalPK).toSet().toList()

val partitionBy = when (val partitionAttr = node.partitionBy) {
is Statement.DDL.PartitionBy.AttrList -> {
val attrListResolved = partitionAttr.attrs.map { attr ->
isDeclaredAttribute(attr, attrs.map { DdlField.fromAttr(it) }) ?: throw IllegalArgumentException("Unresolved ref")
}
statementDDLPartitionByAttrList(attrListResolved.map { Identifier.of(Identifier.Part.delimited(it)) })
}
null -> null
}
return statementDDLCommandCreateTable(
node.name,
attrs,
node.tblConstraints,
node.partitionBy,
partitionBy,
node.tableProperties,
finalPK.map { Identifier.of(Identifier.Part.delimited(it)) },
finalUnique.map { Identifier.of(Identifier.Part.delimited(it)) }
Expand Down Expand Up @@ -1333,6 +1344,7 @@ internal class PlanTyper(private val env: Env, config: Context) {
private fun isDeclaredAttribute(identifier: Identifier, declareAttrs: List<Field>): String? {
declareAttrs.forEach { declared ->
if (identifier.matches(declared.name, identifier.getIdentifier().isRegular()))
// Storing as Declared to work around identifier case sensitivity at the moment
return declared.name
}
return null
Expand Down
2 changes: 1 addition & 1 deletion partiql-planner/src/main/resources/partiql_plan_internal.ion
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ statement::[
],
table_property::{
name: string,
value: partiql_value,
value: string,
},
]
}
Expand Down
Loading

0 comments on commit 14617d0

Please sign in to comment.