github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/doltcore/table/pk_reader.go (about) 1 // Copyright 2020 Dolthub, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package table 16 17 import ( 18 "context" 19 "io" 20 21 "github.com/dolthub/go-mysql-server/sql" 22 23 "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" 24 "github.com/dolthub/dolt/go/libraries/doltcore/row" 25 "github.com/dolthub/dolt/go/libraries/doltcore/schema" 26 "github.com/dolthub/dolt/go/libraries/doltcore/table/typed/noms" 27 "github.com/dolthub/dolt/go/store/types" 28 ) 29 30 type pkTableReader struct { 31 iter types.MapIterator 32 sch schema.Schema 33 } 34 35 var _ SqlTableReader = pkTableReader{} 36 var _ TableReadCloser = pkTableReader{} 37 38 // GetSchema implements the TableReader interface. 39 func (rdr pkTableReader) GetSchema() schema.Schema { 40 return rdr.sch 41 } 42 43 // ReadRow implements the TableReader interface. 44 func (rdr pkTableReader) ReadRow(ctx context.Context) (row.Row, error) { 45 key, val, err := rdr.iter.Next(ctx) 46 47 if err != nil { 48 return nil, err 49 } else if key == nil { 50 return nil, io.EOF 51 } 52 53 return row.FromNoms(rdr.sch, key.(types.Tuple), val.(types.Tuple)) 54 } 55 56 // ReadSqlRow implements the SqlTableReader interface. 57 func (rdr pkTableReader) ReadSqlRow(ctx context.Context) (sql.Row, error) { 58 key, val, err := rdr.iter.Next(ctx) 59 60 if err != nil { 61 return nil, err 62 } else if key == nil { 63 return nil, io.EOF 64 } 65 66 return noms.SqlRowFromTuples(rdr.sch, key.(types.Tuple), val.(types.Tuple)) 67 } 68 69 // Close implements the TableReadCloser interface. 70 func (rdr pkTableReader) Close(_ context.Context) error { 71 return nil 72 } 73 74 func newPkTableReader(ctx context.Context, tbl *doltdb.Table, sch schema.Schema, buffered bool) (pkTableReader, error) { 75 rows, err := tbl.GetRowData(ctx) 76 if err != nil { 77 return pkTableReader{}, err 78 } 79 80 return newPkTableReaderForRows(ctx, rows, sch, buffered) 81 } 82 83 func newPkTableReaderForRows(ctx context.Context, rows types.Map, sch schema.Schema, buffered bool) (pkTableReader, error) { 84 var err error 85 var iter types.MapIterator 86 if buffered { 87 iter, err = rows.Iterator(ctx) 88 } else { 89 iter, err = rows.BufferedIterator(ctx) 90 } 91 if err != nil { 92 return pkTableReader{}, err 93 } 94 95 return pkTableReader{ 96 iter: iter, 97 sch: sch, 98 }, nil 99 } 100 101 func newPkTableReaderFrom(ctx context.Context, tbl *doltdb.Table, sch schema.Schema, val types.Value) (SqlTableReader, error) { 102 rows, err := tbl.GetRowData(ctx) 103 if err != nil { 104 return nil, err 105 } 106 107 iter, err := rows.IteratorFrom(ctx, val) 108 if err != nil { 109 return nil, err 110 } 111 112 return pkTableReader{ 113 iter: iter, 114 sch: sch, 115 }, nil 116 } 117 118 type partitionTableReader struct { 119 SqlTableReader 120 remaining uint64 121 } 122 123 var _ SqlTableReader = &partitionTableReader{} 124 125 func newPkTableReaderForPartition(ctx context.Context, tbl *doltdb.Table, sch schema.Schema, start, end uint64) (SqlTableReader, error) { 126 rows, err := tbl.GetRowData(ctx) 127 if err != nil { 128 return nil, err 129 } 130 131 iter, err := rows.BufferedIteratorAt(ctx, start) 132 if err != nil { 133 return nil, err 134 } 135 136 return &partitionTableReader{ 137 SqlTableReader: pkTableReader{ 138 iter: iter, 139 sch: sch, 140 }, 141 remaining: end - start, 142 }, nil 143 } 144 145 // ReadSqlRow implements the SqlTableReader interface. 146 func (rdr *partitionTableReader) ReadSqlRow(ctx context.Context) (sql.Row, error) { 147 if rdr.remaining == 0 { 148 return nil, io.EOF 149 } 150 rdr.remaining -= 1 151 152 return rdr.SqlTableReader.ReadSqlRow(ctx) 153 }