<commit_before>package main

import (
	"errors"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type storageHandler interface {
	listDatabases() ([]string, error)
	listCollections(dbname string) ([]string, error)
	listMetrics(dbname, collection string) ([]string, error)
	insertSample(dbname, collection string, sample map[string]interface{}) error
	findValues(dbname, collection, metric string) (map[string]float64, error)
	aggregate(dbname, collection, metric string) (map[string]interface{}, error)
	getHeatMap(dbname, collection, metric string) (*heatMap, error)
	getHistogram(dbname, collection, metric string) (map[string]float64, error)
}

type mongoHandler struct {
	Session *mgo.Session
}

func newMongoHandler(addrs []string, timeout time.Duration) (*mongoHandler, error) {
	dialInfo := &mgo.DialInfo{
		Addrs:   addrs,
		Timeout: timeout,
	}

	logger.Info("Connecting to database...")
	if session, err := mgo.DialWithInfo(dialInfo); err != nil {
		logger.Criticalf("Failed to connect to database: %s", err)
		return nil, err
	} else {
		logger.Info("Connection established.")
		session.SetMode(mgo.Monotonic, true)
		return &mongoHandler{session}, nil
	}
}

var dbPrefix = "perf"

func (mongo *mongoHandler) listDatabases() ([]string, error) {
	if err := mongo.Session.Ping(); err != nil {
		mongo.Session.Refresh()
	}
	allDbs, err := mongo.Session.DatabaseNames()
	if err != nil {
		return nil, err
	}

	dbs := []string{}
	for _, db := range allDbs {
		if strings.HasPrefix(db, dbPrefix) {
			dbs = append(dbs, strings.Replace(db, dbPrefix, "", 1))
		}
	}
	return dbs, nil
}

func (mongo *mongoHandler) listCollections(dbname string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_db := session.DB(dbPrefix + dbname)

	allCollections, err := _db.CollectionNames()
	if err != nil {
		return nil, err
	}

	collections := []string{}
	for _, collection := range allCollections {
		if collection != "system.indexes" {
			collections = append(collections, collection)
		}
	}
	return collections, err
}

func (mongo *mongoHandler) listMetrics(dbname, collection string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var metrics []string
	if err := _collection.Find(bson.M{}).Sort("m").Distinct("m", &metrics); err != nil {
		return nil, err
	}
	return metrics, nil
}

func (mongo *mongoHandler) findValues(dbname, collection, metric string) (map[string]float64, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var docs []map[string]interface{}
	if err := _collection.Find(bson.M{"m": metric}).Sort("ts").All(&docs); err != nil {
		return nil, err
	}
	values := map[string]float64{}
	for _, doc := range docs {
		values[doc["ts"].(string)] = doc["v"].(float64)
	}
	return values, nil
}

func (mongo *mongoHandler) insertSample(dbname, collection string, sample map[string]interface{}) error {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	if err := _collection.Insert(sample); err != nil {
		return err
	}
	logger.Infof("Successfully added new sample to %s.%s", dbname, collection)

	for _, key := range []string{"m", "ts", "v"} {
		if err := _collection.EnsureIndexKey(key); err != nil {
			return err
		}
	}
	return nil
}

func calcPercentile(data []float64, p float64) float64 {
	sort.Float64s(data)

	k := float64(len(data)-1) * p
	f := math.Floor(k)
	c := math.Ceil(k)
	if f == c {
		return data[int(k)]
	}
	return data[int(f)]*(c-k) + data[int(c)]*(k-f)
}

const queryLimit = 10000

func (mongo *mongoHandler) aggregate(dbname, collection, metric string) (map[string]interface{}, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	pipe := _collection.Pipe(
		[]bson.M{
			{
				"$match": bson.M{
					"m": metric,
				},
			},
			{
				"$group": bson.M{
					"_id": bson.M{
						"metric": "$m",
					},
					"avg":   bson.M{"$avg": "$v"},
					"min":   bson.M{"$min": "$v"},
					"max":   bson.M{"$max": "$v"},
					"count": bson.M{"$sum": 1},
				},
			},
		},
	)
	summaries := []map[string]interface{}{}
	if err := pipe.All(&summaries); err != nil {
		return nil, err
	}
	if len(summaries) == 0 {
		return map[string]interface{}{}, nil
	}
	summary := summaries[0]
	delete(summary, "_id")<commit_msg>Suppress _id via projection<commit_after>package main

import (
	"errors"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type storageHandler interface {
	listDatabases() ([]string, error)
	listCollections(dbname string) ([]string, error)
	listMetrics(dbname, collection string) ([]string, error)
	insertSample(dbname, collection string, sample map[string]interface{}) error
	findValues(dbname, collection, metric string) (map[string]float64, error)
	aggregate(dbname, collection, metric string) (map[string]interface{}, error)
	getHeatMap(dbname, collection, metric string) (*heatMap, error)
	getHistogram(dbname, collection, metric string) (map[string]float64, error)
}

type mongoHandler struct {
	Session *mgo.Session
}

func newMongoHandler(addrs []string, timeout time.Duration) (*mongoHandler, error) {
	dialInfo := &mgo.DialInfo{
		Addrs:   addrs,
		Timeout: timeout,
	}

	logger.Info("Connecting to database...")
	if session, err := mgo.DialWithInfo(dialInfo); err != nil {
		logger.Criticalf("Failed to connect to database: %s", err)
		return nil, err
	} else {
		logger.Info("Connection established.")
		session.SetMode(mgo.Monotonic, true)
		return &mongoHandler{session}, nil
	}
}

var dbPrefix = "perf"

func (mongo *mongoHandler) listDatabases() ([]string, error) {
	if err := mongo.Session.Ping(); err != nil {
		mongo.Session.Refresh()
	}
	allDbs, err := mongo.Session.DatabaseNames()
	if err != nil {
		return nil, err
	}

	dbs := []string{}
	for _, db := range allDbs {
		if strings.HasPrefix(db, dbPrefix) {
			dbs = append(dbs, strings.Replace(db, dbPrefix, "", 1))
		}
	}
	return dbs, nil
}

func (mongo *mongoHandler) listCollections(dbname string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_db := session.DB(dbPrefix + dbname)

	allCollections, err := _db.CollectionNames()
	if err != nil {
		return nil, err
	}

	collections := []string{}
	for _, collection := range allCollections {
		if collection != "system.indexes" {
			collections = append(collections, collection)
		}
	}
	return collections, err
}

func (mongo *mongoHandler) listMetrics(dbname, collection string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var metrics []string
	if err := _collection.Find(bson.M{}).Sort("m").Distinct("m", &metrics); err != nil {
		return nil, err
	}
	return metrics, nil
}

func (mongo *mongoHandler) findValues(dbname, collection, metric string) (map[string]float64, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var docs []map[string]interface{}
	if err := _collection.Find(bson.M{"m": metric}).Sort("ts").All(&docs); err != nil {
		return nil, err
	}
	values := map[string]float64{}
	for _, doc := range docs {
		values[doc["ts"].(string)] = doc["v"].(float64)
	}
	return values, nil
}

func (mongo *mongoHandler) insertSample(dbname, collection string, sample map[string]interface{}) error {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	if err := _collection.Insert(sample); err != nil {
		return err
	}
	logger.Infof("Successfully added new sample to %s.%s", dbname, collection)

	for _, key := range []string{"m", "ts", "v"} {
		if err := _collection.EnsureIndexKey(key); err != nil {
			return err
		}
	}
	return nil
}

func calcPercentile(data []float64, p float64) float64 {
	sort.Float64s(data)

	k := float64(len(data)-1) * p
	f := math.Floor(k)
	c := math.Ceil(k)
	if f == c {
		return data[int(k)]
	}
	return data[int(f)]*(c-k) + data[int(c)]*(k-f)
}

const queryLimit = 10000

func (mongo *mongoHandler) aggregate(dbname, collection, metric string) (map[string]interface{}, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	pipe := _collection.Pipe(
		[]bson.M{
			{
				"$match": bson.M{
					"m": metric,
				},
			},
			{
				"$group": bson.M{
					"_id":   nil,
					"avg":   bson.M{"$avg": "$v"},
					"min":   bson.M{"$min": "$v"},
					"max":   bson.M{"$max": "$v"},
					"count": bson.M{"$sum": 1},
				},
			},
			{
				"$project": bson.M{
					"_id":   0,
					"avg":   1,
					"min":   1,
					"max":   1,
					"count": 1,
				},
			},
		},
	)
	summaries := []map[string]interface{}{}
	if err := pipe.All(&summaries); err != nil {
		return nil, err
	}
	if len(summaries) == 0 {
		return map[string]interface{}{}, nil
	}
	summary := summaries[0]
{'fraction_non_alphanumeric': 0.17061923583662714, 'fraction_numerical': 0.007575757575757576, 'mean_word_length': 4.250144092219021, 'pattern_counts': {'":': 33, '<': 3, '<?xml version=': 0, '>': 3, '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': '5949616', 'n_tokens_mistral': 3545, 'n_tokens_neox': 3274, 'n_words': 774}