fix(findincsv): memory optimization (#405)

* fix(findincsv): externalize csv as dict caching in uri pkg

* fix(findincsv): cache csv by key

* fix(findincsv): entries by key memory optim

Author: adrienaury

CHANGELOG.md

@@ -14,6 +14,10 @@ Types of changes
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [1.30.2]
+
+- `Fixed` mask `findInCSV` memory usage
+
 ## [1.30.1]
 
 - `Fixed` mask `command` split command line on space protected by quote

pkg/findincsv/findincsv.go

@@ -6,8 +6,6 @@ import (
 	"hash/fnv"
 	"html/template"
 	"sort"
-	"strconv"
-	"strings"
 	tmpl "text/template"
 	"unicode/utf8"
 
@@ -18,16 +16,16 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-type CSVKey struct {
-	filename string
-	lineKey  string
-}
-
 type JaccardCSV struct {
 	csvLine model.Entry
 	lineKey string
 }
 
+type (
+	fileuri string
+	linekey string
+)
+
 type MaskEngine struct {
 	seeder             model.Seeder
 	templateURI        *template.Template
@@ -36,8 +34,7 @@ type MaskEngine struct {
 	temJaccardCSV      *tmlmask.Engine // template to compute key for a csv entry
 	temJaccardEntry    *tmlmask.Engine // template to compute key for json entry
 	expected           string
-	csvAllreadyRead    map[string][]model.Dictionary
-	csvEntryByKey      map[CSVKey][]model.Entry
+	csvEntryByKey      map[fileuri]map[linekey][]model.Entry
 	header             bool
 	sep                rune
 	comment            rune
@@ -115,8 +112,7 @@ func NewMask(conf model.FindInCSVType, seed int64, seeder model.Seeder) (MaskEng
 		temJaccardCSV,
 		temJaccardEntry,
 		expected,
-		map[string][]model.Dictionary{},
-		map[CSVKey][]model.Entry{},
+		map[fileuri]map[linekey][]model.Entry{},
 		conf.Header,
 		sep, comment, conf.FieldsPerRecord, conf.TrimSpace,
 	}, err
@@ -133,10 +129,10 @@ func (me *MaskEngine) Mask(e model.Entry, context ...model.Dictionary) (model.En
 	if err := me.templateURI.Execute(&filenameBuffer, context[0].UnpackUnordered()); err != nil {
 		return nil, err
 	}
-	filename := filenameBuffer.String()
+	filename := fileuri(filenameBuffer.String())
 
 	// Get ExactMatch results
-	exactMatchFinded, exactMatchResult, err := me.ExactMatch(filename, context)
+	exactMatchFinded, exactMatchResult, err := me.exactMatch(filename, context)
 	if err != nil {
 		return nil, err
 	}
@@ -154,7 +150,7 @@ func (me *MaskEngine) Mask(e model.Entry, context ...model.Dictionary) (model.En
 }
 
 // getJaccardMatchResults calculates Jaccard similarity for the given CSV filename and exact match results.
-func (me *MaskEngine) getJaccardMatchResults(filename string, exactMatchResults []model.Entry, context []model.Dictionary) ([]model.Entry, error) {
+func (me *MaskEngine) getJaccardMatchResults(filename fileuri, exactMatchResults []model.Entry, context []model.Dictionary) ([]model.Entry, error) {
 	var jaccardEntryBuffer bytes.Buffer
 	if err := me.temJaccardEntry.Execute(&jaccardEntryBuffer, context[0].UnpackUnordered()); err != nil {
 		return nil, err
@@ -163,19 +159,14 @@ func (me *MaskEngine) getJaccardMatchResults(filename string, exactMatchResults
 
 	// If no exactMatch config
 	if len(exactMatchResults) < 1 {
-		var csvList []model.Dictionary
-		if _, ok := me.csvAllreadyRead[filename]; !ok {
-			var err error
-			csvList, err = me.readCSV(filename)
-			if err != nil {
-				return nil, err
-			}
-		} else {
-			csvList = me.csvAllreadyRead[filename]
+		csvList, err := me.readCSV(filename)
+		if err != nil {
+			return nil, err
 		}
 
 		var records []JaccardCSV
-		for _, record := range csvList {
+		for i := 0; i < csvList.Len(); i++ {
+			record := csvList.Get(i)
 			lineKey, err := me.computeCSVLineKey(record, false)
 			if err != nil {
 				return nil, err
@@ -198,46 +189,45 @@ func (me *MaskEngine) getJaccardMatchResults(filename string, exactMatchResults
 	return sortBySimilarity(jaccardEntryString, records), nil
 }
 
-func (me *MaskEngine) ExactMatch(filename string, context []model.Dictionary) (bool, []model.Entry, error) {
+func (me *MaskEngine) exactMatch(filename fileuri, context []model.Dictionary) (bool, []model.Entry, error) {
 	if me.temExactMatchEntry != nil && me.temExactMatchCSV != nil {
-		var csvList []model.Dictionary
-		if _, ok := me.csvAllreadyRead[filename]; !ok {
-			var err error
-			csvList, err = me.readCSV(filename)
-			if err != nil {
-				return false, nil, err
-			}
+		csvList, err := me.readCSV(filename)
+		if err != nil {
+			return false, nil, err
 		}
 
 		var exactEntryBuffer bytes.Buffer
 		if err := me.temExactMatchEntry.Execute(&exactEntryBuffer, context[0].UnpackUnordered()); err != nil {
 			return false, nil, err
 		}
-		exactEntryString := exactEntryBuffer.String()
-		err := me.getExactMatchCsvResult(filename, csvList)
+		exactEntryString := linekey(exactEntryBuffer.String())
+		err = me.getExactMatchCsvResult(filename, csvList)
 		if err != nil {
 			return false, []model.Entry{}, err
 		}
 
-		results := me.csvEntryByKey[CSVKey{
-			filename: filename,
-			lineKey:  exactEntryString,
-		}]
-		if len(results) < 1 {
-			return false, results, nil
-		}
-		return true, results, nil
+		results := me.readCsvEntryByKey(filename, exactEntryString)
+
+		return len(results) > 0, results, nil
 	}
 	return true, []model.Entry{}, nil
 }
 
-func (me *MaskEngine) readCSV(filename string) ([]model.Dictionary, error) {
-	recordsFromFile, err := uri.ReadCsv(filename, me.sep, me.comment, me.fieldsPerRecord, me.trimSpaces)
+func (me *MaskEngine) readCsvEntryByKey(filename fileuri, exactEntryString linekey) []model.Entry {
+	cache, cacheExists := me.csvEntryByKey[filename]
+	if !cacheExists {
+		panic("csv file is not cached, please report the bug on GitHub CGI-FR")
+	}
+
+	return cache[exactEntryString]
+}
+
+func (me *MaskEngine) readCSV(filename fileuri) (uri.DictRecords, error) {
+	recordsFromFile, err := uri.ReadCsvAsDicts(string(filename), me.sep, me.comment, me.fieldsPerRecord, me.trimSpaces, me.header)
 	if err != nil {
 		return nil, err
 	}
-	csvList := me.createEntriesFromCSVLines(recordsFromFile)
-	return csvList, nil
+	return recordsFromFile, nil
 }
 
 func (me *MaskEngine) computeCSVLineKey(record model.Dictionary, exactMatch bool) (string, error) {
@@ -258,60 +248,32 @@ func (me *MaskEngine) computeCSVLineKey(record model.Dictionary, exactMatch bool
 	return output.String(), nil
 }
 
-func (me *MaskEngine) getExactMatchCsvResult(filename string, csvList []model.Dictionary) error {
-	for _, record := range csvList {
-		lineKey, err := me.computeCSVLineKey(record, true)
-		if err != nil {
-			return err
-		}
+func (me *MaskEngine) getExactMatchCsvResult(filename fileuri, csvList uri.DictRecords) error {
+	_, cacheExists := me.csvEntryByKey[filename]
+	if !cacheExists {
+		cache := map[linekey][]model.Entry{}
 
-		key := CSVKey{
-			filename: filename,
-			lineKey:  lineKey,
-		}
+		for i := 0; i < csvList.Len(); i++ {
+			record := csvList.Get(i)
+			lineKey, err := me.computeCSVLineKey(record, true)
+			if err != nil {
+				return err
+			}
 
-		if records, ok := me.csvEntryByKey[key]; ok {
-			records = append(records, record)
-			me.csvEntryByKey[key] = records
-		} else {
-			me.csvEntryByKey[key] = []model.Entry{record}
+			if records, ok := cache[linekey(lineKey)]; ok {
+				records = append(records, record)
+				cache[linekey(lineKey)] = records
+			} else {
+				cache[linekey(lineKey)] = []model.Entry{record}
+			}
 		}
+
+		me.csvEntryByKey[filename] = cache
 	}
 
 	return nil
 }
 
-func (me *MaskEngine) createEntriesFromCSVLines(records uri.CSVRecords) []model.Dictionary {
-	results := []model.Dictionary{}
-
-	for i := 0; i < records.Len(); i++ {
-		record := records.Get(i)
-		if me.header {
-			obj := model.NewDictionary()
-			headers := records.Get(0)
-			for i, header := range headers {
-				if me.trimSpaces {
-					obj.Set(strings.TrimSpace(header), strings.TrimSpace(record[i]))
-				} else {
-					obj.Set(header, record[i])
-				}
-			}
-			results = append(results, obj)
-		} else {
-			obj := model.NewDictionary()
-			for i, value := range record {
-				if me.trimSpaces {
-					obj.Set(strconv.Itoa(i), strings.TrimSpace(value))
-				} else {
-					obj.Set(strconv.Itoa(i), value)
-				}
-			}
-			results = append(results, obj)
-		}
-	}
-	return results
-}
-
 // Get numbers of result waited in expected config, by default return as at-least-one
 func (me *MaskEngine) getExpectedResult(results []model.Entry) (model.Entry, error) {
 	resultCount := len(results)
@@ -337,7 +299,7 @@ func (me *MaskEngine) getExpectedResult(results []model.Entry) (model.Entry, err
 }
 
 // JaccardSimilarity calculates the Jaccard similarity between two strings.
-func JaccardSimilarity(s1, s2 string) float64 {
+func jaccardSimilarity(s1, s2 string) float64 {
 	if s1 == s2 {
 		return 1.0
 	}
@@ -398,7 +360,7 @@ func sortBySimilarity(jaccardEntryString string, list []JaccardCSV) []model.Entr
 	var entriesWithSimilarity []EntryWithSimilarity
 
 	for _, record := range list {
-		similarity := JaccardSimilarity(jaccardEntryString, record.lineKey)
+		similarity := jaccardSimilarity(jaccardEntryString, record.lineKey)
 		entriesWithSimilarity = append(entriesWithSimilarity, EntryWithSimilarity{Key: record.lineKey, Entry: record.csvLine, Similarity: similarity})
 	}
 

pkg/findincsv/findincsv_test.go

@@ -465,7 +465,7 @@ func TestJaccardSimilarityShouldReturnWaitedNumber(t *testing.T) {
 	entry := "1 Boulevard Albert Einstain"
 	csv := "1 Boulevard Albert Einstein"
 
-	waited := JaccardSimilarity(entry, csv)
+	waited := jaccardSimilarity(entry, csv)
 	expect := 0.8518518518518519
 	assert.Equal(t, expect, waited)
 }

pkg/uri/uri.go

@@ -26,6 +26,7 @@ import (
 	"net/url"
 	"os"
 	"strconv"
+	"strings"
 
 	"github.com/cgi-fr/pimo/pkg/maskingdata"
 	"github.com/cgi-fr/pimo/pkg/model"
@@ -42,6 +43,7 @@ type Records[T model.Entry | []string] interface {
 type (
 	CSVRecords   Records[[]string]
 	EntryRecords Records[model.Entry]
+	DictRecords  Records[model.Dictionary]
 )
 
 type records[T model.Entry | []string] struct {
@@ -63,10 +65,69 @@ func (r records[T]) Collect() []T {
 }
 
 var (
-	cacheCSV   map[string]records[[]string]    = map[string]records[[]string]{}
-	cacheEntry map[string]records[model.Entry] = map[string]records[model.Entry]{}
+	cacheCSV   map[string]records[[]string]         = map[string]records[[]string]{}
+	cacheEntry map[string]records[model.Entry]      = map[string]records[model.Entry]{}
+	cacheDict  map[string]records[model.Dictionary] = map[string]records[model.Dictionary]{}
 )
 
+func ReadCsvAsDicts(uri string, sep rune, comment rune, fieldsPerRecord int, trimLeadingSpaces bool, hasHeaders bool) (DictRecords, error) {
+	if records, present := cacheDict[uri]; present {
+		return records, nil
+	}
+
+	csvRecords, err := ReadCsv(uri, sep, comment, fieldsPerRecord, trimLeadingSpaces)
+	if err != nil {
+		return nil, err
+	}
+
+	headers := []string{}
+	if hasHeaders {
+		headerRecord := csvRecords.Get(0)
+		for _, header := range headerRecord {
+			if trimLeadingSpaces {
+				headers = append(headers, strings.TrimSpace(header))
+			} else {
+				headers = append(headers, header)
+			}
+		}
+	}
+
+	i := 0
+	if hasHeaders {
+		i = 1
+	}
+
+	results := make([]model.Dictionary, csvRecords.Len())
+	for ; i < csvRecords.Len(); i++ {
+		record := csvRecords.Get(i)
+		obj := model.NewDictionary()
+		if hasHeaders {
+			for i, header := range headers {
+				if trimLeadingSpaces {
+					obj.Set(header, strings.TrimSpace(record[i]))
+				} else {
+					obj.Set(header, record[i])
+				}
+			}
+		} else {
+			for i, value := range record {
+				if trimLeadingSpaces {
+					obj.Set(strconv.Itoa(i), strings.TrimSpace(value))
+				} else {
+					obj.Set(strconv.Itoa(i), value)
+				}
+			}
+		}
+		results[i] = obj
+	}
+
+	if hasHeaders {
+		results = results[1:]
+	}
+
+	return records[model.Dictionary]{results}, nil
+}
+
 func ReadCsv(uri string, sep rune, comment rune, fieldsPerRecord int, trimLeadingSpaces bool) (CSVRecords, error) {
 	if records, present := cacheCSV[uri]; present {
 		return records, nil