Skip to content

Commit

Permalink
support SELECT FROM TABLESAMPLE syntax (pingcap#1071)
Browse files Browse the repository at this point in the history
* support SELECT FROM TABLESAMPLE syntax

* refine naming

* make table sample clause as a part of TableFactor

* address comment
  • Loading branch information
tangenta committed Nov 10, 2020
1 parent 0b08570 commit 9881cb2
Show file tree
Hide file tree
Showing 6 changed files with 8,931 additions and 8,546 deletions.
12 changes: 7 additions & 5 deletions ast/ddl.go
Original file line number Diff line number Diff line change
Expand Up @@ -330,13 +330,15 @@ func (n *ReferenceDef) Accept(v Visitor) (Node, bool) {
return v.Leave(newNode)
}
n = newNode.(*ReferenceDef)
node, ok := n.Table.Accept(v)
if !ok {
return n, false
if n.Table != nil {
node, ok := n.Table.Accept(v)
if !ok {
return n, false
}
n.Table = node.(*TableName)
}
n.Table = node.(*TableName)
for i, val := range n.IndexPartSpecifications {
node, ok = val.Accept(v)
node, ok := val.Accept(v)
if !ok {
return n, false
}
Expand Down
111 changes: 109 additions & 2 deletions ast/dml.go
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,7 @@ type TableName struct {

IndexHints []*IndexHint
PartitionNames []model.CIStr
TableSample *TableSample
}

// Restore implements Node interface.
Expand Down Expand Up @@ -218,14 +219,22 @@ func (n *TableName) restoreIndexHints(ctx *format.RestoreCtx) error {
return errors.Annotate(err, "An error occurred while splicing IndexHints")
}
}

return nil
}

func (n *TableName) Restore(ctx *format.RestoreCtx) error {
n.restoreName(ctx)
n.restorePartitions(ctx)
return n.restoreIndexHints(ctx)
if err := n.restoreIndexHints(ctx); err != nil {
return err
}
if n.TableSample != nil {
ctx.WritePlain(" ")
if err := n.TableSample.Restore(ctx); err != nil {
return errors.Annotate(err, "An error occurred while splicing TableName.TableSample")
}
}
return nil
}

// IndexHintType is the type for index hint use, ignore or force.
Expand Down Expand Up @@ -304,6 +313,13 @@ func (n *TableName) Accept(v Visitor) (Node, bool) {
return v.Leave(newNode)
}
n = newNode.(*TableName)
if n.TableSample != nil {
newTs, ok := n.TableSample.Accept(v)
if !ok {
return n, false
}
n.TableSample = newTs.(*TableSample)
}
return v.Leave(n)
}

Expand Down Expand Up @@ -411,6 +427,12 @@ func (n *TableSource) Restore(ctx *format.RestoreCtx) error {
if err := tn.restoreIndexHints(ctx); err != nil {
return errors.Annotate(err, "An error occurred while restore TableSource.Source.(*TableName).IndexHints")
}
if tn.TableSample != nil {
ctx.WritePlain(" ")
if err := tn.TableSample.Restore(ctx); err != nil {
return errors.Annotate(err, "An error occurred while splicing TableName.TableSample")
}
}

if needParen {
ctx.WritePlain(")")
Expand Down Expand Up @@ -777,6 +799,90 @@ func (n *OrderByClause) Accept(v Visitor) (Node, bool) {
return v.Leave(n)
}

type SampleMethodType int8

const (
SampleMethodTypeNone SampleMethodType = iota
SampleMethodTypeSystem
SampleMethodTypeBernoulli
SampleMethodTypeTiDBRegion
)

type SampleClauseUnitType int8

const (
SampleClauseUnitTypeDefault SampleClauseUnitType = iota
SampleClauseUnitTypeRow
SampleClauseUnitTypePercent
)

type TableSample struct {
node
SampleMethod SampleMethodType
Expr ExprNode
SampleClauseUnit SampleClauseUnitType
RepeatableSeed ExprNode
}

func (s *TableSample) Restore(ctx *format.RestoreCtx) error {
ctx.WriteKeyWord("TABLESAMPLE ")
switch s.SampleMethod {
case SampleMethodTypeBernoulli:
ctx.WriteKeyWord("BERNOULLI ")
case SampleMethodTypeSystem:
ctx.WriteKeyWord("SYSTEM ")
case SampleMethodTypeTiDBRegion:
ctx.WriteKeyWord("REGION ")
}
ctx.WritePlain("(")
if s.Expr != nil {
if err := s.Expr.Restore(ctx); err != nil {
return errors.Annotate(err, "An error occurred while restore TableSample.Expr")
}
}
switch s.SampleClauseUnit {
case SampleClauseUnitTypeDefault:
case SampleClauseUnitTypePercent:
ctx.WriteKeyWord(" PERCENT")
case SampleClauseUnitTypeRow:
ctx.WriteKeyWord(" ROWS")

}
ctx.WritePlain(")")
if s.RepeatableSeed != nil {
ctx.WriteKeyWord(" REPEATABLE")
ctx.WritePlain("(")
if err := s.RepeatableSeed.Restore(ctx); err != nil {
return errors.Annotate(err, "An error occurred while restore TableSample.Expr")
}
ctx.WritePlain(")")
}
return nil
}

func (s *TableSample) Accept(v Visitor) (node Node, ok bool) {
newNode, skipChildren := v.Enter(s)
if skipChildren {
return v.Leave(newNode)
}
s = newNode.(*TableSample)
if s.Expr != nil {
node, ok = s.Expr.Accept(v)
if !ok {
return s, false
}
s.Expr = node.(ExprNode)
}
if s.RepeatableSeed != nil {
node, ok = s.RepeatableSeed.Accept(v)
if !ok {
return s, false
}
s.RepeatableSeed = node.(ExprNode)
}
return v.Leave(s)
}

type SelectStmtKind uint8

const (
Expand Down Expand Up @@ -911,6 +1017,7 @@ func (n *SelectStmt) Restore(ctx *format.RestoreCtx) error {
if n.From == nil && n.Where != nil {
ctx.WriteKeyWord(" FROM DUAL")
}

if n.Where != nil {
ctx.WriteKeyWord(" WHERE ")
if err := n.Where.Restore(ctx); err != nil {
Expand Down
4 changes: 4 additions & 0 deletions misc.go
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,7 @@ var tokenMap = map[string]int{
"BACKUPS": backups,
"BEGIN": begin,
"BETWEEN": between,
"BERNOULLI": bernoulli,
"BIGINT": bigIntType,
"BINARY": binaryType,
"BINDING": binding,
Expand Down Expand Up @@ -508,6 +509,7 @@ var tokenMap = map[string]int{
"PARTITIONING": partitioning,
"PARTITIONS": partitions,
"PASSWORD": password,
"PERCENT": percent,
"PER_DB": per_db,
"PER_TABLE": per_table,
"PESSIMISTIC": pessimistic,
Expand Down Expand Up @@ -657,10 +659,12 @@ var tokenMap = map[string]int{
"SUPER": super,
"SWAPS": swaps,
"SWITCHES": switchesSym,
"SYSTEM": system,
"SYSTEM_TIME": systemTime,
"TABLE_CHECKSUM": tableChecksum,
"TABLE": tableKwd,
"TABLES": tables,
"TABLESAMPLE": tableSample,
"TABLESPACE": tablespace,
"TELEMETRY": telemetry,
"TELEMETRY_ID": telemetryID,
Expand Down

0 comments on commit 9881cb2

Please sign in to comment.