Skip to content

Commit

Permalink
Add support for create table statement in planner
Browse files Browse the repository at this point in the history
  • Loading branch information
@yliuuuu
yliuuuu committed Dec 26, 2024
1 parent ea779ee commit 434a601
Show file tree
Hide file tree
Showing 23 changed files with 2,132 additions and 132 deletions.
4 changes: 2 additions & 2 deletions partiql-ast/src/main/java/org/partiql/ast/DataType.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ public static class StructField extends AstNode {

public final boolean isOptional;

@Nullable
@NotNull
public final List<AttributeConstraint> constraints;
@Nullable
public final String comment;
Expand All @@ -43,7 +43,7 @@ public StructField(
@NotNull Identifier name,
@NotNull DataType type,
boolean isOptional,
@Nullable List<AttributeConstraint> constraints,
@NotNull List<AttributeConstraint> constraints,
@Nullable String comment) {
this.name = name;
this.type = type;
Expand Down
7 changes: 7 additions & 0 deletions partiql-plan/src/main/java/org/partiql/plan/Action.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import org.jetbrains.annotations.NotNull;
import org.partiql.plan.rex.Rex;
import org.partiql.spi.catalog.Table;

/**
* A PartiQL statement action within a plan.
Expand All @@ -19,4 +20,10 @@ public interface Query extends Action {
@NotNull
public Rex getRex();
}

// A better way to segment is to have an object interface in SPI
public interface CreateTable extends Action {
@NotNull
public Table getTable();
}
}
51 changes: 51 additions & 0 deletions partiql-planner/src/main/kotlin/org/partiql/planner/internal/DdlField.kt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
package org.partiql.planner.internal

import org.partiql.planner.internal.ir.Statement
import org.partiql.planner.internal.ir.statementDDLAttribute
import org.partiql.planner.internal.typer.CompilerType
import org.partiql.spi.catalog.Identifier
import org.partiql.types.Field
import org.partiql.types.PType
import org.partiql.types.shape.PShape

/**
* An implementation for [Field] that is used by DDL
* to hold additional information in the struct field.
* It is identical to [Statement.DDL.Attribute]
*/
internal data class DdlField(
val name: Identifier,
val type: PShape,
val isNullable: Boolean,
val isOptional: Boolean,
val constraints: List<Statement.DDL.Constraint>,
val isPrimaryKey: Boolean,
val isUnique: Boolean,
val comment: String?
) : Field {

override fun getName(): String {
return name.getIdentifier().getText()
}

override fun getType(): PType {
return type
}

fun toAttr() =
statementDDLAttribute(
this.name,
this.type,
this.isNullable,
this.isOptional,
this.isPrimaryKey,
this.isUnique,
this.constraints,
this.comment
)

companion object {
fun fromAttr(attr: Statement.DDL.Attribute): DdlField =
DdlField(attr.name, attr.type, attr.isNullable, attr.isOptional, attr.constraints, attr.isPrimaryKey, attr.isUnique, attr.comment)
}
}
171 changes: 171 additions & 0 deletions 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 @@ -62,6 +62,12 @@ import org.partiql.planner.internal.ir.builder.RexOpTupleUnionBuilder
import org.partiql.planner.internal.ir.builder.RexOpVarGlobalBuilder
import org.partiql.planner.internal.ir.builder.RexOpVarLocalBuilder
import org.partiql.planner.internal.ir.builder.RexOpVarUnresolvedBuilder
import org.partiql.planner.internal.ir.builder.StatementDdlAttributeBuilder
import org.partiql.planner.internal.ir.builder.StatementDdlBuilder
import org.partiql.planner.internal.ir.builder.StatementDdlCommandCreateTableBuilder
import org.partiql.planner.internal.ir.builder.StatementDdlConstraintCheckBuilder
import org.partiql.planner.internal.ir.builder.StatementDdlPartitionByAttrListBuilder
import org.partiql.planner.internal.ir.builder.StatementDdlTablePropertyBuilder
import org.partiql.planner.internal.ir.builder.StatementQueryBuilder
import org.partiql.planner.internal.ir.visitor.PlanVisitor
import org.partiql.planner.internal.typer.CompilerType
Expand All @@ -73,6 +79,7 @@ import org.partiql.spi.function.Function
import org.partiql.value.PartiQLValue
import org.partiql.value.PartiQLValueExperimental
import kotlin.random.Random
import org.partiql.types.shape.PShape

internal abstract class PlanNode {
@JvmField
Expand Down Expand Up @@ -177,6 +184,7 @@ internal sealed class Ref : PlanNode() {
internal sealed class Statement : PlanNode() {
public override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R = when (this) {
is Query -> visitor.visitStatementQuery(this, ctx)
is DDL -> visitor.visitStatementDDL(this, ctx)
}

internal data class Query(
Expand All @@ -195,6 +203,169 @@ internal sealed class Statement : PlanNode() {
internal fun builder(): StatementQueryBuilder = StatementQueryBuilder()
}
}

internal data class DDL(
@JvmField
internal val command: Command,
) : Statement() {
override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.add(command)
kids.filterNotNull()
}


override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitStatementDDL(this, ctx)

sealed class Command : PlanNode() {
override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R = when (this) {
is CreateTable -> visitor.visitStatementDDLCommandCreateTable(this, ctx)
}

internal data class CreateTable(
@JvmField
internal val name: Identifier,
@JvmField
internal val attributes: List<Attribute>,
@JvmField
internal val tblConstraints: List<Constraint>,
@JvmField
internal val partitionBy: PartitionBy?,
@JvmField
internal val tableProperties: List<TableProperty>,
@JvmField
internal val primaryKey: List<Identifier>,
@JvmField
internal val unique: List<Identifier>,
) : Command() {
override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.addAll(attributes)
kids.addAll(tblConstraints)
partitionBy?.let { kids.add(it) }
kids.addAll(tableProperties)
kids.filterNotNull()
}


override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitStatementDDLCommandCreateTable(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): StatementDdlCommandCreateTableBuilder =
StatementDdlCommandCreateTableBuilder()
}
}
}

internal data class Attribute(
@JvmField
internal val name: Identifier,
@JvmField
internal val type: PShape,
@JvmField
internal val isNullable: Boolean,
@JvmField
internal val isOptional: Boolean,
@JvmField
internal val isPrimaryKey: Boolean,
@JvmField
internal val isUnique: Boolean,
@JvmField
internal val constraints: List<Constraint>,
@JvmField
val comment: String?
) : PlanNode() {
override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.addAll(constraints)
kids.filterNotNull()
}

override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitStatementDDLAttribute(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): StatementDdlAttributeBuilder = StatementDdlAttributeBuilder()
}
}

internal sealed class Constraint : PlanNode() {
override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R = when (this) {
is Check -> visitor.visitStatementDDLConstraintCheck(this, ctx)
}

internal data class Check(
@JvmField
internal val expression: Rex,
@JvmField
val sql: String
) : Constraint() {
override val children: List<PlanNode> by lazy {
val kids = mutableListOf<PlanNode?>()
kids.add(expression)
kids.filterNotNull()
}


override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitStatementDDLConstraintCheck(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): StatementDdlConstraintCheckBuilder =
StatementDdlConstraintCheckBuilder()
}
}
}

internal sealed class PartitionBy : PlanNode() {
override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R = when (this) {
is AttrList -> visitor.visitStatementDDLPartitionByAttrList(this, ctx)
}

internal data class AttrList(
@JvmField
internal val attrs: List<Identifier>,
) : PartitionBy() {
override val children: List<PlanNode> = emptyList()

override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitStatementDDLPartitionByAttrList(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): StatementDdlPartitionByAttrListBuilder =
StatementDdlPartitionByAttrListBuilder()
}
}
}

internal data class TableProperty(
@JvmField
internal val name: String,
@JvmField
internal val `value`: String,
) : PlanNode() {
override val children: List<PlanNode> = emptyList()

override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
visitor.visitStatementDDLTableProperty(this, ctx)

internal companion object {
@JvmStatic
internal fun builder(): StatementDdlTablePropertyBuilder = StatementDdlTablePropertyBuilder()
}
}

internal companion object {
@JvmStatic
internal fun builder(): StatementDdlBuilder = StatementDdlBuilder()
}
}
}

internal data class Rex(
Expand Down
6 changes: 5 additions & 1 deletion ...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 @@ -23,6 +23,7 @@ import org.partiql.ast.FromExpr
import org.partiql.ast.FromJoin
import org.partiql.ast.FromTableRef
import org.partiql.ast.FromType
import org.partiql.ast.Query
import org.partiql.ast.QueryBody
import org.partiql.ast.Statement
import org.partiql.ast.expr.Expr
Expand All @@ -33,7 +34,10 @@ import org.partiql.planner.internal.helpers.toBinder
*/
internal object NormalizeFromSource : AstPass {

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

private object Visitor : AstRewriter<Int>() {

Expand Down
6 changes: 5 additions & 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 @@ -19,6 +19,7 @@ import org.partiql.ast.Ast.groupByKey
import org.partiql.ast.AstNode
import org.partiql.ast.AstRewriter
import org.partiql.ast.GroupBy
import org.partiql.ast.Query
import org.partiql.ast.Statement
import org.partiql.ast.expr.Expr
import org.partiql.planner.internal.helpers.toBinder
Expand All @@ -28,7 +29,10 @@ import org.partiql.planner.internal.helpers.toBinder
*/
internal object NormalizeGroupBy : AstPass {

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

private object Visitor : AstRewriter<Int>() {

Expand Down
Loading

0 comments on commit 434a601

Please sign in to comment.