-
Notifications
You must be signed in to change notification settings - Fork 23
/
union_all.go
140 lines (119 loc) · 3.69 KB
/
union_all.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
package exec
import (
"github.com/squareup/pranadb/common"
)
type UnionAll struct {
pushExecutorBase
genIDColIndex int
}
func NewUnionAll() (*UnionAll, error) {
return &UnionAll{
pushExecutorBase: pushExecutorBase{},
}, nil
}
func (u *UnionAll) ReCalcSchemaFromChildren() error {
child0 := u.children[0]
u.colTypes = child0.ColTypes()
u.colsVisible = child0.ColsVisible()
u.genIDColIndex = len(child0.ColTypes())
u.colTypes = append(u.colTypes, common.VarcharColumnType)
u.colsVisible = append(u.colsVisible, false)
u.keyCols = []int{u.genIDColIndex}
u.rowsFactory = common.NewRowsFactory(u.colTypes)
// When a row is incoming we need to know which source it came from as we generate the new key from this, in order to
// do this we create an intermediate object for each input which calls HandleRowsWithIndex passing in the index
var newChildren []PushExecutor
for i, exec := range u.children {
forwarder := &RowWithIndexForwarder{
pushExecutorBase: pushExecutorBase{},
index: i,
unionAll: u,
}
newChildren = append(newChildren, forwarder)
exec.SetParent(forwarder)
forwarder.SetParent(u)
forwarder.AddChild(exec)
}
u.children = newChildren
return nil
}
func (u *UnionAll) HandleRows(rowsBatch RowsBatch, ctx *ExecutionContext) error {
panic("should not be called")
}
func (u *UnionAll) HandleRowsWithIndex(index int, rowsBatch RowsBatch, ctx *ExecutionContext) error {
numRows := rowsBatch.Len()
// TODO this could be optimised by just taking the columns from the incoming chunk and adding them to the new rows
out := u.rowsFactory.NewRows(numRows)
entries := make([]RowsEntry, numRows)
rc := 0
for i := 0; i < numRows; i++ {
prevRow := rowsBatch.PreviousRow(i)
pi := -1
if prevRow != nil {
if err := u.appendRow(index, prevRow, out); err != nil {
return err
}
pi = rc
rc++
}
currRow := rowsBatch.CurrentRow(i)
ci := -1
if currRow != nil {
if err := u.appendRow(index, currRow, out); err != nil {
return err
}
ci = rc
rc++
}
entries[i] = NewRowsEntry(pi, ci, rowsBatch.ReceiverIndex(i))
}
return u.parent.HandleRows(NewRowsBatch(out, entries), ctx)
}
func (u *UnionAll) appendRow(index int, inRow *common.Row, out *common.Rows) error {
for i := 0; i < len(u.colTypes)-1; i++ {
if inRow.IsNull(i) {
out.AppendNullToColumn(i)
} else {
colType := u.colTypes[i]
switch colType.Type {
case common.TypeTinyInt, common.TypeInt, common.TypeBigInt:
out.AppendInt64ToColumn(i, inRow.GetInt64(i))
case common.TypeDouble:
out.AppendFloat64ToColumn(i, inRow.GetFloat64(i))
case common.TypeVarchar:
out.AppendStringToColumn(i, inRow.GetString(i))
case common.TypeTimestamp:
out.AppendTimestampToColumn(i, inRow.GetTimestamp(i))
case common.TypeDecimal:
out.AppendDecimalToColumn(i, inRow.GetDecimal(i))
default:
panic("unexpected column type")
}
}
}
id, err := u.generateID(inRow, index)
if err != nil {
return err
}
out.AppendStringToColumn(u.genIDColIndex, id)
return nil
}
func (u *UnionAll) generateID(row *common.Row, index int) (string, error) {
actualChild := u.children[index].GetChildren()[0]
// We create the key by prefixing with the index, then appending the key from the actual child
var key []byte
key = common.AppendUint32ToBufferLE(key, uint32(index))
key, err := common.EncodeKeyCols(row, actualChild.KeyCols(), actualChild.ColTypes(), key)
if err != nil {
return "", err
}
return string(key), nil
}
type RowWithIndexForwarder struct {
pushExecutorBase
index int
unionAll *UnionAll
}
func (r *RowWithIndexForwarder) HandleRows(rowsBatch RowsBatch, ctx *ExecutionContext) error {
return r.unionAll.HandleRowsWithIndex(r.index, rowsBatch, ctx)
}