<commit_before>package memory

import (
	"github.com/xephonhq/xephon-k/pkg/common"
)

// Index is a map of inverted index with tag name as key and tag value as term for the inverted index
type Index struct {
	tagKeyIndex     map[string]map[string]bool // map[string]bool is used as set
	invertedIndexes map[string]*InvertedIndex
}

// InvertedIndex use Term for tag value postings for a list of sorted series ID
// TODO: Series ID should use locality sensitive hashing https://en.wikipedia.org/wiki/Locality-sensitive_hashing
type InvertedIndex struct {
	Term     string
	Postings []common.SeriesID
}

var initialPostingSize = 10

func NewIndex(capacity int) *Index {
	return &Index{
		tagKeyIndex:     make(map[string]map[string]bool, capacity),
		invertedIndexes: make(map[string]*InvertedIndex, capacity),
	}
}

func newInvertedIndex(term string) *InvertedIndex {
	return &InvertedIndex{
		Term:     term,
		Postings: make([]common.SeriesID, 0, initialPostingSize),
	}
}

func (idx *Index) Filter(f *common.Filter) []common.SeriesID {
	// TODO: we need locking to ensure the correctness
	// can we have multiple read lock on a same object?
	empty := []common.SeriesID{}
	switch f.Type {
	case "tag_match":
		// TODO: what if the tag user provided is invalid, the should be checked at API
		// and we assume everything is right here?
		return idx.Get(f.Key, f.Value)
	case "and":
		return Intersect(idx.Filter(f.LeftOperand), idx.Filter(f.RightOperand))
	case "or":
		return Union(idx.Filter(f.LeftOperand), idx.Filter(f.RightOperand))
	case "in":
		// TODO: in is just multiple and?
		log.Warn("in is not implemented")
		return empty
	default:
		// TODO: this should be checked in upper level
		log.Warn("%s is unsupported", f.Type)
		return empty
	}
}

func (idx *Index) Get(tagKey string, tagValue string) []common.SeriesID {
	term := Term(tagKey, tagValue)
	iidx, ok := idx.invertedIndexes[term]
	if ok {
		return iidx.Postings
	} else {
		return []common.SeriesID{}
	}
}

func (idx *Index) Add(id common.SeriesID, tagKey string, tagValue string) {
	// update tagKeyIndex
	_, ok := idx.tagKeyIndex[tagKey]
	if !ok {
		idx.tagKeyIndex[tagKey] = make(map[string]bool)
	}
	idx.tagKeyIndex[tagKey][tagValue] = true

	// TODO: should add separator, in Prometheus `db.go` it's `const sep = '\xff'`
	term := Term(tagKey, tagValue)
	// create the inverted index if not exists
	_, ok = idx.invertedIndexes[term]
	if !ok {
		idx.invertedIndexes[term] = newInvertedIndex(term)
	}
	idx.invertedIndexes[term].Add(id)
}

// TODO: actually we can have a fixed size map to cache the hot series, so there is no need to lookup if the id is already in there
func (iidx *InvertedIndex) Add(id common.SeriesID) {
	// binary search and insert the value if not found
	low, high := 0, len(iidx.Postings)
	for low < high {
		// TODO: use custom compare function or compare it directly using <
		mid := low + (high-low)/2 // avoid overflow, copied from `src/sort/search.go` sort.Search
		if iidx.Postings[mid] >= id {<commit_msg>[mem] Add mutex for in memory index as well<commit_after>package memory

import (
	"github.com/xephonhq/xephon-k/pkg/common"
	"sync"
)

// Index is a map of inverted index with tag name as key and tag value as term for the inverted index
type Index struct {
	mu              sync.RWMutex
	tagKeyIndex     map[string]map[string]bool // map[string]bool is used as set
	invertedIndexes map[string]*InvertedIndex
}

// InvertedIndex use Term for tag value postings for a list of sorted series ID
// TODO: Series ID should use locality sensitive hashing https://en.wikipedia.org/wiki/Locality-sensitive_hashing
type InvertedIndex struct {
	mu       sync.RWMutex
	Term     string
	Postings []common.SeriesID
}

var initialPostingSize = 10

func NewIndex(capacity int) *Index {
	return &Index{
		tagKeyIndex:     make(map[string]map[string]bool, capacity),
		invertedIndexes: make(map[string]*InvertedIndex, capacity),
	}
}

func newInvertedIndex(term string) *InvertedIndex {
	return &InvertedIndex{
		Term:     term,
		Postings: make([]common.SeriesID, 0, initialPostingSize),
	}
}

func (idx *Index) Filter(f *common.Filter) []common.SeriesID {
	// TODO: we need locking to ensure the correctness
	// can we have multiple read lock on a same object?
	empty := []common.SeriesID{}
	switch f.Type {
	case "tag_match":
		// TODO: what if the tag user provided is invalid, the should be checked at API
		// and we assume everything is right here?
		return idx.Get(f.Key, f.Value)
	case "and":
		return Intersect(idx.Filter(f.LeftOperand), idx.Filter(f.RightOperand))
	case "or":
		return Union(idx.Filter(f.LeftOperand), idx.Filter(f.RightOperand))
	case "in":
		// TODO: in is just multiple and?
		log.Warn("in is not implemented")
		return empty
	default:
		// TODO: this should be checked in upper level
		log.Warn("%s is unsupported", f.Type)
		return empty
	}
}

func (idx *Index) Get(tagKey string, tagValue string) []common.SeriesID {
	term := Term(tagKey, tagValue)
	iidx, ok := idx.invertedIndexes[term]
	if ok {
		return iidx.Postings
	} else {
		return []common.SeriesID{}
	}
}

func (idx *Index) Add(id common.SeriesID, tagKey string, tagValue string) {
	idx.mu.Lock()
	// update tagKeyIndex
	_, ok := idx.tagKeyIndex[tagKey]
	if !ok {
		idx.tagKeyIndex[tagKey] = make(map[string]bool)
	}
	idx.tagKeyIndex[tagKey][tagValue] = true

	// TODO: should add separator, in Prometheus `db.go` it's `const sep = '\xff'`
	term := Term(tagKey, tagValue)
	// create the inverted index if not exists
	_, ok = idx.invertedIndexes[term]
	if !ok {
		idx.invertedIndexes[term] = newInvertedIndex(term)
	}
	// NOTE: we unlock here because each inverted index also have its own lock
	idx.mu.Unlock()
	idx.invertedIndexes[term].Add(id)
}

// TODO: actually we can have a fixed size map to cache the hot series, so there is no need to lookup if the id is already in there
func (iidx *InvertedIndex) Add(id common.SeriesID) {
	iidx.mu.Lock()
	defer iidx.mu.Unlock()
	// binary search and insert the value if not found
	low, high := 0, len(iidx.Postings)
	for low < high {
		// TODO: use custom compare function or compare it directly using <
		mid := low + (high-low)/2 // avoid overflow, copied from `src/sort/search.go` sort.Search
		if iidx.Postings[mid] >= id {
{'fraction_non_alphanumeric': 0.1136327313408982, 'fraction_numerical': 0.00159821000479463, 'mean_word_length': 3.97456279809221, 'pattern_counts': {'":': 8, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 2, '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': '330933', 'n_tokens_mistral': 2132, 'n_tokens_neox': 2065, 'n_words': 685}