<commit_before>/* Coordination between a collection of chunks. */
package db

import (
	"errors"
	"fmt"
	"github.com/HouzuoGuo/tiedot/chunk"
	"github.com/HouzuoGuo/tiedot/chunkfile"
	"github.com/HouzuoGuo/tiedot/tdlog"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"sync"
)

type Col struct {
	BaseDir       string // Collection dir path
	Chunks        []*chunk.ChunkCol
	ChunkMutexes  []*sync.RWMutex // Synchronize access to chunks
	NewChunkMutex sync.Mutex      // Synchronize creation of new chunk
	NumChunks     uint64          // Total number of chunks
}

// Open a collection (made of chunks).
func OpenCol(baseDir string) (col *Col, err error) {
	// Create the directory if it does not yet exist
	if err = os.MkdirAll(baseDir, 0700); err != nil {
		return
	}
	col = &Col{BaseDir: baseDir, NewChunkMutex: sync.Mutex{},
		Chunks: make([]*chunk.ChunkCol, 0), ChunkMutexes: make([]*sync.RWMutex, 0)}
	// Walk the collection directory and look for how many chunks there are
	maxChunk := 0
	walker := func(currPath string, info os.FileInfo, err2 error) error {
		if err2 != nil {
			// log and skip the error
			tdlog.Error(err)
			return nil
		}
		if info.IsDir() {
			// According to the directory name - chunk number, look for how many chunks there are
			if chunkNum, err := strconv.Atoi(info.Name()); err == nil {
				if chunkNum > maxChunk {
					maxChunk = chunkNum
				}
			}
			return nil
		}
		return nil
	}
	if err = filepath.Walk(baseDir, walker); err != nil {
		return
	}
	col.NumChunks = uint64(maxChunk) + 1
	// Open all chunks - there is at least one chunk
	tdlog.Printf("Opening %d chunks in collection %s", maxChunk+1, baseDir)
	for i := 0; i <= maxChunk; i++ {
		var oneChunk *chunk.ChunkCol
		oneChunk, err = chunk.OpenChunk(uint64(i), path.Join(baseDir, strconv.Itoa(i)))
		col.Chunks = append(col.Chunks, oneChunk)
		col.ChunkMutexes = append(col.ChunkMutexes, new(sync.RWMutex))
	}
	return
}

// Create a new chunk.
func (col *Col) CreateNewChunk() {
	col.NewChunkMutex.Lock()
	tdlog.Printf("Going to create a new chunk (number %d) in collection %s", col.NumChunks, col.BaseDir)
	newChunk, err := chunk.OpenChunk(col.NumChunks, path.Join(col.BaseDir, strconv.Itoa(int(col.NumChunks))))
	if err != nil {
		col.NewChunkMutex.Unlock()
		return
	}
	// Make indexes
	for _, path := range col.Chunks[0].HTPaths {
		newChunk.Index(path)
	}
	// Put the new chunk into col structures
	col.Chunks = append(col.Chunks, newChunk)
	col.ChunkMutexes = append(col.ChunkMutexes, new(sync.RWMutex))
	col.NumChunks += 1
	col.NewChunkMutex.Unlock()
}

// Insert a new document, return new document's ID.
func (col *Col) Insert(doc interface{}) (id uint64, err error) {
	// Try to insert the doc into a random chunk
	randChunkNum := uint64(rand.Int63n(int64(col.NumChunks)))
	randChunk := col.Chunks[randChunkNum]<commit_msg>better error handling in chunk creation<commit_after>/* Coordination between a collection of chunks. */
package db

import (
	"errors"
	"fmt"
	"github.com/HouzuoGuo/tiedot/chunk"
	"github.com/HouzuoGuo/tiedot/chunkfile"
	"github.com/HouzuoGuo/tiedot/tdlog"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"sync"
)

type Col struct {
	BaseDir       string // Collection dir path
	Chunks        []*chunk.ChunkCol
	ChunkMutexes  []*sync.RWMutex // Synchronize access to chunks
	NewChunkMutex sync.Mutex      // Synchronize creation of new chunk
	NumChunks     uint64          // Total number of chunks
}

// Open a collection (made of chunks).
func OpenCol(baseDir string) (col *Col, err error) {
	// Create the directory if it does not yet exist
	if err = os.MkdirAll(baseDir, 0700); err != nil {
		return
	}
	col = &Col{BaseDir: baseDir, NewChunkMutex: sync.Mutex{},
		Chunks: make([]*chunk.ChunkCol, 0), ChunkMutexes: make([]*sync.RWMutex, 0)}
	// Walk the collection directory and look for how many chunks there are
	maxChunk := 0
	walker := func(currPath string, info os.FileInfo, err2 error) error {
		if err2 != nil {
			// log and skip the error
			tdlog.Error(err)
			return nil
		}
		if info.IsDir() {
			// According to the directory name - chunk number, look for how many chunks there are
			if chunkNum, err := strconv.Atoi(info.Name()); err == nil {
				if chunkNum > maxChunk {
					maxChunk = chunkNum
				}
			}
			return nil
		}
		return nil
	}
	if err = filepath.Walk(baseDir, walker); err != nil {
		return
	}
	col.NumChunks = uint64(maxChunk) + 1
	// Open all chunks - there is at least one chunk
	tdlog.Printf("Opening %d chunks in collection %s", maxChunk+1, baseDir)
	for i := 0; i <= maxChunk; i++ {
		var oneChunk *chunk.ChunkCol
		oneChunk, err = chunk.OpenChunk(uint64(i), path.Join(baseDir, strconv.Itoa(i)))
		col.Chunks = append(col.Chunks, oneChunk)
		col.ChunkMutexes = append(col.ChunkMutexes, new(sync.RWMutex))
	}
	return
}

// Create a new chunk.
func (col *Col) CreateNewChunk() {
	col.NewChunkMutex.Lock()
	defer col.NewChunkMutex.Unlock()
	tdlog.Printf("Going to create a new chunk (number %d) in collection %s", col.NumChunks, col.BaseDir)
	newChunk, err := chunk.OpenChunk(col.NumChunks, path.Join(col.BaseDir, strconv.Itoa(int(col.NumChunks))))
	if err != nil {
		col.NewChunkMutex.Unlock()
		return
	}
	// Make indexes
	for _, path := range col.Chunks[0].HTPaths {
		if path[0] != chunk.UID_PATH {
			if err := newChunk.Index(path); err != nil {
				tdlog.Panicf("Failed to create index %s, error: %v", path, err)
			}
		}
	}
	// Put the new chunk into col structures
	col.Chunks = append(col.Chunks, newChunk)
	col.ChunkMutexes = append(col.ChunkMutexes, new(sync.RWMutex))
	col.NumChunks += 1
}

// Insert a new document, return new document's ID.
func (col *Col) Insert(doc interface{}) (id uint64, err error) {
	// Try to insert the doc into a random chunk
	randChunkNum := uint64(rand.Int63n(int64(col.NumChunks)))
	randChunk := col.Chunks[randChunkNum]
{'fraction_non_alphanumeric': 0.12809988028048572, 'fraction_numerical': 0.009748589020010261, 'mean_word_length': 3.604724409448819, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2420119', 'n_tokens_mistral': 2187, 'n_tokens_neox': 2105, 'n_words': 595}