internal/index/merge.go

   1 package index
   2
   3 import (
   4         "bufio"
   5         "encoding/binary"
   6         "io"
   7         "log"
   8         "os"
   9         "path/filepath"
  10         "sort"
  11 )
  12
  13 type fileMetaEntry struct {
  14         idxid   uint32
  15         entries uint32
  16         offset  int64
  17 }
  18
  19 type indexMeta struct {
  20         docidBase uint32
  21         rd        *PForReader
  22 }
  23
  24 type posrelMetaEntry struct {
  25         idxid  uint32
  26         offset int64
  27 }
  28
  29 type posrelMeta struct {
  30         rd *PosrelReader
  31 }
  32
  33 func readMeta(dir, typ string, idx map[Trigram][]uint32, idxid uint32) error {
  34         f, err := os.Open(filepath.Join(dir, "posting."+typ+".meta"))
  35         if err != nil {
  36                 return err
  37         }
  38         defer f.Close()
  39         st, err := f.Stat()
  40         if err != nil {
  41                 return err
  42         }
  43         bufr := bufio.NewReader(f)
  44
  45         buf := make([]byte, metaEntrySize)
  46         for i := 0; i < (int(st.Size()) / metaEntrySize); i++ {
  47                 if _, err := io.ReadFull(bufr, buf); err != nil {
  48                         return err
  49                 }
  50                 t := Trigram(binary.LittleEndian.Uint32(buf))
  51                 idx[t] = append(idx[t], idxid)
  52         }
  53         return nil
  54 }
  55
  56 func readPosrelMeta(dir string, idx map[Trigram][]posrelMetaEntry, idxid uint32) error {
  57         f, err := os.Open(filepath.Join(dir, "posting.posrel.meta"))
  58         if err != nil {
  59                 return err
  60         }
  61         defer f.Close()
  62         st, err := f.Stat()
  63         if err != nil {
  64                 return err
  65         }
  66         bufr := bufio.NewReader(f)
  67
  68         buf := make([]byte, metaEntrySize)
  69         var entry MetaEntry
  70         for i := 0; i < (int(st.Size()) / metaEntrySize); i++ {
  71                 if _, err := io.ReadFull(bufr, buf); err != nil {
  72                         return err
  73                 }
  74                 entry.Unmarshal(buf)
  75                 idx[entry.Trigram] = append(idx[entry.Trigram], posrelMetaEntry{
  76                         idxid:  idxid,
  77                         offset: entry.OffsetData,
  78                 })
  79         }
  80         return nil
  81 }
  82
  83 const debug = false
  84
  85 var debugTrigram = func(trigram string) Trigram {
  86         t := []byte(trigram)
  87         return Trigram(uint32(t[0])<<16 | uint32(t[1])<<8 | uint32(t[2]))
  88 }("_op")
  89
  90 func ConcatN(destdir string, srcdirs []string) error {
  91         fDocidMap, err := os.Create(filepath.Join(destdir, "docid.map"))
  92         if err != nil {
  93                 return err
  94         }
  95         defer fDocidMap.Close()
  96         cw := newCountingWriter(fDocidMap)
  97
  98         var (
  99                 base    uint32
 100                 offsets []uint32
 101         )
 102         bufr := bufio.NewReader(nil)
 103         bases := make([]uint32, len(srcdirs))
 104         for idx, dir := range srcdirs {
 105                 bases[idx] = base
 106                 f, err := os.Open(filepath.Join(dir, "docid.map"))
 107                 if err != nil {
 108                         return err
 109                 }
 110                 defer f.Close()
 111                 st, err := f.Stat()
 112                 if err != nil {
 113                         return err
 114                 }
 115                 if _, err := f.Seek(-4, io.SeekEnd); err != nil {
 116                         return err
 117                 }
 118                 // Locate index offset:
 119                 var indexOffset uint32
 120                 if err := binary.Read(f, binary.LittleEndian, &indexOffset); err != nil {
 121                         return err
 122                 }
 123
 124                 // TODO: detect |base| overflows
 125                 n := (uint32(st.Size()) - indexOffset - 4) / 4
 126                 log.Printf("%s (idx %d) contains %d docids", dir, idx, n)
 127                 base += n
 128
 129                 if _, err := f.Seek(0, io.SeekStart); err != nil {
 130                         return err
 131                 }
 132                 bufr.Reset(f)
 133                 // TODO(performance): measure whether using the index and incrementing
 134                 // the offsets is any faster than this method:
 135                 scanner := bufio.NewScanner(&io.LimitedReader{
 136                         R: bufr,
 137                         N: int64(indexOffset)})
 138                 for scanner.Scan() {
 139                         offsets = append(offsets, uint32(cw.offset))
 140                         cw.Write(scanner.Bytes())
 141                         cw.Write([]byte{'\n'})
 142                 }
 143                 if err := scanner.Err(); err != nil {
 144                         return err
 145                 }
 146         }
 147         indexStart := uint32(cw.offset)
 148         if err := binary.Write(&cw, binary.LittleEndian, offsets); err != nil {
 149                 return err
 150         }
 151         if err := binary.Write(&cw, binary.LittleEndian, indexStart); err != nil {
 152                 return err
 153         }
 154
 155         if err := cw.Close(); err != nil {
 156                 return err
 157         }
 158
 159         log.Printf("reading fileMetaEntries")
 160
 161         idxMetaDocid := make([]indexMeta, len(srcdirs))
 162         idxMetaPos := make([]indexMeta, len(srcdirs))
 163         idxMetaPosrel := make([]posrelMeta, len(srcdirs))
 164
 165         idxDocid := make(map[Trigram][]uint32)
 166         for idx, dir := range srcdirs {
 167                 base := bases[idx]
 168
 169                 {
 170                         rd, err := newPForReader(dir, "docid")
 171                         if err != nil {
 172                                 return err
 173                         }
 174                         defer rd.Close()
 175                         idxMetaDocid[idx] = indexMeta{docidBase: base, rd: rd}
 176                 }
 177
 178                 if err := readMeta(dir, "docid", idxDocid, uint32(idx)); err != nil {
 179                         return err
 180                 }
 181
 182                 {
 183                         rd, err := newPForReader(dir, "pos")
 184                         if err != nil {
 185                                 return err
 186                         }
 187                         defer rd.Close()
 188                         idxMetaPos[idx] = indexMeta{docidBase: base, rd: rd}
 189                 }
 190
 191                 {
 192                         rd, err := newPosrelReader(dir)
 193                         if err != nil {
 194                                 return err
 195                         }
 196                         defer rd.Close()
 197
 198                         idxMetaPosrel[idx] = posrelMeta{rd: rd}
 199                 }
 200         }
 201
 202         trigrams := make([]Trigram, 0, len(idxDocid))
 203         for t := range idxDocid {
 204                 trigrams = append(trigrams, t)
 205         }
 206         sort.Slice(trigrams, func(i, j int) bool { return trigrams[i] < trigrams[j] })
 207
 208         {
 209                 log.Printf("writing merged docids")
 210                 dw, err := newPForWriter(destdir, "docid")
 211                 if err != nil {
 212                         return err
 213                 }
 214
 215                 fDocidMeta, err := os.Create(filepath.Join(destdir, "posting.docid.meta"))
 216                 if err != nil {
 217                         return err
 218                 }
 219                 defer fDocidMeta.Close()
 220                 bufwDocidMeta := bufio.NewWriter(fDocidMeta)
 221
 222                 meBuf := make([]byte, metaEntrySize)
 223                 dr := NewDeltaReader()
 224                 for _, t := range trigrams {
 225                         if debug {
 226                                 if t != debugTrigram {
 227                                         continue
 228                                 }
 229                         }
 230
 231                         //for _, t := range []trigram{trigram(6650227), trigram(7959906)} {
 232                         //ctrl, data := dw.Offsets()
 233                         me := MetaEntry{
 234                                 Trigram: t,
 235                                 //OffsetCtrl: ctrl,
 236                                 //OffsetEnc:  data,
 237                                 OffsetData: dw.Offset(),
 238                         }
 239                         var last uint32
 240                         for _, idxid := range idxDocid[t] {
 241                                 idx := idxMetaDocid[idxid]
 242                                 meta, err := idx.rd.metaEntry1(t)
 243                                 if err != nil {
 244                                         if err == errNotFound {
 245                                                 continue
 246                                         }
 247                                         return err
 248                                 }
 249                                 me.Entries += meta.Entries
 250                                 dr.Reset(meta, idx.rd.data.Data)
 251                                 docids := dr.Read() // returns non-nil at least once
 252                                 // Bump the first docid: it needs to be mapped from the old
 253                                 // docid range [0, n) to the new docid range [base, base+n).
 254                                 //
 255                                 // Since we are building a single docid list for this trigram,
 256                                 // the new value needs to be a delta, hence, subtract last.
 257                                 docids[0] += (idx.docidBase - last)
 258                                 for docids != nil {
 259                                         for _, d := range docids {
 260                                                 if err := dw.PutUint32(d); err != nil {
 261                                                         return err
 262                                                 }
 263                                                 last += d
 264                                         }
 265                                         docids = dr.Read()
 266                                 }
 267                         }
 268                         if err := dw.Flush(); err != nil {
 269                                 return err
 270                         }
 271                         me.Marshal(meBuf)
 272                         if _, err := bufwDocidMeta.Write(meBuf); err != nil {
 273                                 //if err := binary.Write(bufwDocidMeta, binary.LittleEndian, &me); err != nil {
 274                                 return err
 275                         }
 276                 }
 277
 278                 if err := bufwDocidMeta.Flush(); err != nil {
 279                         return err
 280                 }
 281
 282                 if err := fDocidMeta.Close(); err != nil {
 283                         return err
 284                 }
 285
 286                 if err := dw.Close(); err != nil {
 287                         return err
 288                 }
 289         }
 290
 291         {
 292                 log.Printf("writing merged posrel")
 293                 fmeta, err := os.Create(filepath.Join(destdir, "posting.posrel.meta"))
 294                 if err != nil {
 295                         return err
 296                 }
 297                 defer fmeta.Close()
 298                 bufwmeta := bufio.NewWriter(fmeta)
 299
 300                 fposrel, err := os.Create(filepath.Join(destdir, "posting.posrel.data"))
 301                 if err != nil {
 302                         return err
 303                 }
 304                 defer fposrel.Close()
 305                 cw := newCountingWriter(fposrel)
 306                 pw := newPosrelWriter(&cw)
 307                 for _, t := range trigrams {
 308                         if debug {
 309                                 if t != debugTrigram {
 310                                         continue
 311                                 }
 312                         }
 313                         if t == 2105376 { // TODO: document: "   "?
 314                                 continue
 315                         }
 316
 317                         me := MetaEntry{
 318                                 Trigram:    t,
 319                                 OffsetData: int64(cw.offset),
 320                         }
 321                         if err := binary.Write(bufwmeta, binary.LittleEndian, &me); err != nil {
 322                                 return err
 323                         }
 324                         for _, idxid := range idxDocid[t] {
 325                                 // TODO: refactor all metaEntry1 to use ,ok idiom, they only ever return errNotFound
 326                                 fmeta, err := idxMetaPos[idxid].rd.metaEntry1(t)
 327                                 if err != nil {
 328                                         if err == errNotFound {
 329                                                 continue
 330                                         }
 331                                         return err
 332                                 }
 333
 334                                 pmeta, err := idxMetaPosrel[idxid].rd.metaEntry1(t)
 335                                 if err != nil {
 336                                         if err == errNotFound {
 337                                                 continue
 338                                         }
 339                                         return err
 340                                 }
 341                                 b := idxMetaPosrel[idxid].rd.data.Data[pmeta.OffsetData:]
 342                                 if err := pw.Write(b, int(fmeta.Entries)); err != nil {
 343                                         return err
 344                                 }
 345
 346                         }
 347                         if err := pw.Flush(); err != nil {
 348                                 return err
 349                         }
 350                 }
 351                 if err := bufwmeta.Flush(); err != nil {
 352                         return err
 353                 }
 354                 if err := fmeta.Close(); err != nil {
 355                         return err
 356                 }
 357                 if err := cw.Close(); err != nil {
 358                         return err
 359                 }
 360         }
 361
 362         {
 363                 log.Printf("writing merged pos")
 364                 dw, err := newPForWriter(destdir, "pos")
 365                 if err != nil {
 366                         return err
 367                 }
 368
 369                 fDocidMeta, err := os.Create(filepath.Join(destdir, "posting.pos.meta"))
 370                 if err != nil {
 371                         return err
 372                 }
 373                 defer fDocidMeta.Close()
 374                 bufwDocidMeta := bufio.NewWriter(fDocidMeta)
 375
 376                 meBuf := make([]byte, metaEntrySize)
 377                 dr := NewDeltaReader()
 378                 //for _, t := range []trigram{trigram(6650227), trigram(7959906)} {
 379                 for _, t := range trigrams {
 380                         if debug {
 381                                 if t != debugTrigram {
 382                                         continue
 383                                 }
 384                         }
 385
 386                         if t == 2105376 { // TODO: document: "   "?
 387                                 continue
 388                         }
 389
 390                         //ctrl, data := dw.Offsets()
 391                         me := MetaEntry{
 392                                 Trigram: t,
 393                                 // OffsetCtrl: ctrl,
 394                                 // OffsetEnc:  data,
 395                                 OffsetData: dw.Offset(),
 396                         }
 397
 398                         for _, idxid := range idxDocid[t] {
 399                                 idx := idxMetaPos[idxid]
 400                                 meta, err := idx.rd.metaEntry1(t)
 401                                 if err != nil {
 402                                         if err == errNotFound {
 403                                                 continue
 404                                         }
 405                                         return err
 406                                 }
 407                                 me.Entries += meta.Entries
 408                                 dr.Reset(meta, idx.rd.data.Data)
 409
 410                                 for docids := dr.Read(); docids != nil; docids = dr.Read() {
 411                                         for _, d := range docids {
 412                                                 if err := dw.PutUint32(d); err != nil {
 413                                                         return err
 414                                                 }
 415                                         }
 416                                 }
 417                         }
 418
 419                         if err := dw.Flush(); err != nil {
 420                                 return err
 421                         }
 422
 423                         me.Marshal(meBuf)
 424                         if _, err := bufwDocidMeta.Write(meBuf); err != nil {
 425                                 //if err := binary.Write(bufwDocidMeta, binary.LittleEndian, &me); err != nil {
 426                                 return err
 427                         }
 428                 }
 429
 430                 if err := bufwDocidMeta.Flush(); err != nil {
 431                         return err
 432                 }
 433
 434                 if err := fDocidMeta.Close(); err != nil {
 435                         return err
 436                 }
 437
 438                 if err := dw.Close(); err != nil {
 439                         return err
 440                 }
 441         }
 442
 443         return nil
 444 }