<commit_before>package join

import (
	"fmt"
	"sort"
	"strings"

	radix "github.com/armon/go-radix"
	"github.com/jacksontj/dataman/src/storage_node/metadata"
)

// TODO: map[string]Filter
type JoinMap map[string]map[string]interface{}

func ParseJoinMap(joinField interface{}) (JoinMap, error) {
	joinMap := make(JoinMap)
	switch joinFieldTyped := joinField.(type) {
	case []interface{}:
		for _, item := range joinFieldTyped {
			stringKey, ok := item.(string)
			if !ok {
				return nil, fmt.Errorf("joinFieldList must be []string")
			}
			joinMap[stringKey] = make(map[string]interface{})
		}

	case []string:
		for _, item := range joinFieldTyped {
			joinMap[item] = make(map[string]interface{})
		}

	case map[string]interface{}:
		for k, v := range joinFieldTyped {
			if vTyped, ok := v.(map[string]interface{}); ok {
				joinMap[k] = vTyped
			} else {
				return nil, fmt.Errorf("values must be map[string]interface{}")
			}
		}

	default:
		// TODO: be more explicit
		return nil, fmt.Errorf("Invalid join type")
	}
	return joinMap, nil
}

func SortJoinFieldList(l []string) {
	// Sort joinFieldList shortest -> longest -- so we do the deps first
	less := func(i, j int) bool {
		return strings.Count(l[i], ".") < strings.Count(l[j], ".")
	}
	sort.Slice(l, less)
}

// Repsonsible for splitting and normalizing, not sorting
func SplitJoinFieldsString(collectionName string, joinMap JoinMap) ([]string, []string) {
	NormalizeJoinMap(collectionName, joinMap)
	forwardList := make([]string, 0)
	reverseList := make([]string, 0)
	for joinItem, _ := range joinMap {
		if strings.HasPrefix(joinItem, ".") {
			forwardList = append(forwardList, joinItem)
		} else {
			reverseList = append(reverseList, joinItem)
		}
	}
	SortJoinFieldList(forwardList)
	SortJoinFieldList(reverseList)

	return forwardList, reverseList
}

func NormalizeJoinMap(n string, m JoinMap) {
	prefix := n + "."
	for k, v := range m {
		if strings.HasPrefix(k, prefix) {
			m[strings.TrimPrefix(k, n)] = v
			delete(m, k)
		}
	}
}

// Since router and storage_node have different collections we need an interface which is common)
type MetaCollection interface {
	GetName() string
	GetFieldByName(name string) *metadata.CollectionField
	GetField(nameParts []string) *metadata.CollectionField
}

type CollectionGetter func(name string) MetaCollection

func NewCollection(m MetaCollection) *Collection {
	return &Collection{
		M:           m,
		Name:        m.GetName(),
		ForwardJoin: make([]*ForwardJoin, 0),
	}
}

// Represent a collection in the join state
type Collection struct {
	// A pointer to us-- so we can do things like find fields etc.
	M    MetaCollection `json:"-"`
	Name string         `json:"name"`

	ForwardJoin []*ForwardJoin `json:"forward_join,omitempty"`
	// There can be N reverse ones
	ReverseJoin []*ReverseJoin `json:"reverse_join,omitempty"`
}

func (c *Collection) HasJoins() bool {
	return !(len(c.ForwardJoin) == 0 && len(c.ReverseJoin) == 0)
}

type ForwardJoin struct {
	Key       string                    `json:"key"`
	C         *Collection               `json:"c"`
	JoinField *metadata.CollectionField `json:"join_field"`
	Filter    map[string]interface{}    `json:"filter"`
}

type ReverseJoin struct {
	Key       string                    `json:"key"`
	C         *Collection               `json:"c"`
	JoinField *metadata.CollectionField `json:"join_field"`
	Filter    map[string]interface{}    `json:"filter"`
}

// Return a list of join steps. These can be run in chains
func OrderJoins(collectionMetaGetter CollectionGetter, collection MetaCollection, joinMap JoinMap) (*Collection, error) {
	// Normalize all the ".<key>" to "<collection>.<key>"
	NormalizeJoinMap(collection.GetName(), joinMap)

	forwardJoinList, reverseJoinList := SplitJoinFieldsString(collection.GetName(), joinMap)

	// Map of collection -> joinCollection repr -- this is for reverse joins (basically)
	thisCollection := NewCollection(collection)
	collectionMap := map[string][]*Collection{collection.GetName(): []*Collection{thisCollection}}

	forwardJoinFunc := func(thisCollection *Collection, forwardJoinList []string) error {
		// For forward joins we want to check for prefix matching and do our mapping that way
		prefixTree := radix.New()

		// Do forward joins
		for _, join := range forwardJoinList {
			joinParts := strings.Split(join, ".")
			joinParts = joinParts[1:]

			var joinField *metadata.CollectionField
			baseCollection := thisCollection
			if prefix, item, ok := prefixTree.LongestPrefix(join); ok {
				baseJoinCollection := item.(*Collection)
				baseCollection = baseJoinCollection

				suffix := strings.TrimPrefix(join, prefix+".")
				joinParts = joinParts[strings.Count(prefix, "."):]

				joinField = baseJoinCollection.M.GetFieldByName(suffix)
			} else {
				joinField = thisCollection.M.GetField(joinParts)
			}

			// TODO: check for non-existance
			metaJoinCollection := collectionMetaGetter(joinField.Relation.Collection)
			if metaJoinCollection == nil {
				return fmt.Errorf("Unknown collection: %s", joinField.Relation.Collection)
			}
			joinCollection := NewCollection(metaJoinCollection)

			forwardJoin := &ForwardJoin{
				Key:       strings.Join(joinParts, "."),
				C:         joinCollection,
				JoinField: joinField,
				Filter:    joinMap[join],
			}
<commit_msg>Error earlier for invalid joinfield<commit_after>package join

import (
	"fmt"
	"sort"
	"strings"

	radix "github.com/armon/go-radix"
	"github.com/jacksontj/dataman/src/storage_node/metadata"
)

// TODO: map[string]Filter
type JoinMap map[string]map[string]interface{}

func ParseJoinMap(joinField interface{}) (JoinMap, error) {
	joinMap := make(JoinMap)
	switch joinFieldTyped := joinField.(type) {
	case []interface{}:
		for _, item := range joinFieldTyped {
			stringKey, ok := item.(string)
			if !ok {
				return nil, fmt.Errorf("joinFieldList must be []string")
			}
			joinMap[stringKey] = make(map[string]interface{})
		}

	case []string:
		for _, item := range joinFieldTyped {
			joinMap[item] = make(map[string]interface{})
		}

	case map[string]interface{}:
		for k, v := range joinFieldTyped {
			if vTyped, ok := v.(map[string]interface{}); ok {
				joinMap[k] = vTyped
			} else {
				return nil, fmt.Errorf("values must be map[string]interface{}")
			}
		}

	default:
		// TODO: be more explicit
		return nil, fmt.Errorf("Invalid join type")
	}
	return joinMap, nil
}

func SortJoinFieldList(l []string) {
	// Sort joinFieldList shortest -> longest -- so we do the deps first
	less := func(i, j int) bool {
		return strings.Count(l[i], ".") < strings.Count(l[j], ".")
	}
	sort.Slice(l, less)
}

// Repsonsible for splitting and normalizing, not sorting
func SplitJoinFieldsString(collectionName string, joinMap JoinMap) ([]string, []string) {
	NormalizeJoinMap(collectionName, joinMap)
	forwardList := make([]string, 0)
	reverseList := make([]string, 0)
	for joinItem, _ := range joinMap {
		if strings.HasPrefix(joinItem, ".") {
			forwardList = append(forwardList, joinItem)
		} else {
			reverseList = append(reverseList, joinItem)
		}
	}
	SortJoinFieldList(forwardList)
	SortJoinFieldList(reverseList)

	return forwardList, reverseList
}

func NormalizeJoinMap(n string, m JoinMap) {
	prefix := n + "."
	for k, v := range m {
		if strings.HasPrefix(k, prefix) {
			m[strings.TrimPrefix(k, n)] = v
			delete(m, k)
		}
	}
}

// Since router and storage_node have different collections we need an interface which is common)
type MetaCollection interface {
	GetName() string
	GetFieldByName(name string) *metadata.CollectionField
	GetField(nameParts []string) *metadata.CollectionField
}

type CollectionGetter func(name string) MetaCollection

func NewCollection(m MetaCollection) *Collection {
	return &Collection{
		M:           m,
		Name:        m.GetName(),
		ForwardJoin: make([]*ForwardJoin, 0),
	}
}

// Represent a collection in the join state
type Collection struct {
	// A pointer to us-- so we can do things like find fields etc.
	M    MetaCollection `json:"-"`
	Name string         `json:"name"`

	ForwardJoin []*ForwardJoin `json:"forward_join,omitempty"`
	// There can be N reverse ones
	ReverseJoin []*ReverseJoin `json:"reverse_join,omitempty"`
}

func (c *Collection) HasJoins() bool {
	return !(len(c.ForwardJoin) == 0 && len(c.ReverseJoin) == 0)
}

type ForwardJoin struct {
	Key       string                    `json:"key"`
	C         *Collection               `json:"c"`
	JoinField *metadata.CollectionField `json:"join_field"`
	Filter    map[string]interface{}    `json:"filter"`
}

type ReverseJoin struct {
	Key       string                    `json:"key"`
	C         *Collection               `json:"c"`
	JoinField *metadata.CollectionField `json:"join_field"`
	Filter    map[string]interface{}    `json:"filter"`
}

// Return a list of join steps. These can be run in chains
func OrderJoins(collectionMetaGetter CollectionGetter, collection MetaCollection, joinMap JoinMap) (*Collection, error) {
	// Normalize all the ".<key>" to "<collection>.<key>"
	NormalizeJoinMap(collection.GetName(), joinMap)

	forwardJoinList, reverseJoinList := SplitJoinFieldsString(collection.GetName(), joinMap)

	// Map of collection -> joinCollection repr -- this is for reverse joins (basically)
	thisCollection := NewCollection(collection)
	collectionMap := map[string][]*Collection{collection.GetName(): []*Collection{thisCollection}}

	forwardJoinFunc := func(thisCollection *Collection, forwardJoinList []string) error {
		// For forward joins we want to check for prefix matching and do our mapping that way
		prefixTree := radix.New()

		// Do forward joins
		for _, join := range forwardJoinList {
			joinParts := strings.Split(join, ".")
			joinParts = joinParts[1:]

			var joinField *metadata.CollectionField
			baseCollection := thisCollection
			if prefix, item, ok := prefixTree.LongestPrefix(join); ok {
				baseJoinCollection := item.(*Collection)
				baseCollection = baseJoinCollection

				suffix := strings.TrimPrefix(join, prefix+".")
				joinParts = joinParts[strings.Count(prefix, "."):]

				joinField = baseJoinCollection.M.GetFieldByName(suffix)
			} else {
				joinField = thisCollection.M.GetField(joinParts)
			}

			if joinField == nil {
				return fmt.Errorf("Invalid join key: %s", join)
			}

			// TODO: check for non-existance
			metaJoinCollection := collectionMetaGetter(joinField.Relation.Collection)
			if metaJoinCollection == nil {
				return fmt.Errorf("Unknown collection: %s", joinField.Relation.Collection)
			}
			joinCollection := NewCollection(metaJoinCollection)

			forwardJoin := &ForwardJoin{
				Key:       strings.Join(joinParts, "."),
				C:         joinCollection,
				JoinField: joinField,
				Filter:    joinMap[join],
			}

{'fraction_non_alphanumeric': 0.12529175613854915, 'fraction_numerical': 0.0011203435720287555, 'mean_word_length': 3.9432395016151363, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 13, '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': '6486249', 'n_tokens_mistral': 3537, 'n_tokens_neox': 3239, 'n_words': 867}