join.go

package simplecsv

import (
	"errors"
	"fmt"
	"strings"
)

// JoinByField returns a new csv with the inner join of s and other by the
// column fieldName, which must exist in both csvs. Only the rows whose join
// value exists in both csvs are in the result. If the same join value shows
// up more than once, all combinations of rows are in the result.
//
// Join values are compared exactly: case matters ("Ana" != "ana").
// Rows shorter than the join column position are treated as having an empty
// string in that cell, and cells beyond the header width are not in the
// result.
//
// The result headers are the headers of s followed by the headers of other
// excluding the join column. Non-join columns of other must not share names
// with columns of s; a collision is rejected so header names stay unique.
//
// If fieldName does not exist in one of the csvs, or if the result would have
// duplicate headers, it returns a copy of s and false. The source csvs are not
// modified and the result shares no rows with them.
func (s SimpleCsv) JoinByField(other SimpleCsv, fieldName string) (SimpleCsv, bool) {
	return s.joinAll(other, fieldName, fieldName, false, false)
}

// LeftJoinByField is like JoinByField, but the rows of s without a match in
// other are also in the result, with empty strings in the columns of other.
func (s SimpleCsv) LeftJoinByField(other SimpleCsv, fieldName string) (SimpleCsv, bool) {
	return s.joinAll(other, fieldName, fieldName, true, false)
}

// Join returns a new csv with the inner join of s and other by the column
// leftField of s and the column rightField of other, which may have different
// names. Only the rows whose join value exists in both csvs are in the
// result. If the same join value shows up more than once, all combinations
// of rows are in the result.
//
// Join values are compared exactly: case matters ("Ana" != "ana").
// Rows shorter than the join column position are treated as having an empty
// string in that cell, and cells beyond the header width are not in the
// result.
//
// The result headers are the headers of s followed by the headers of other
// excluding the join column rightField (the left join column is kept). Non-join
// columns of other must not share names with columns of s; a collision is
// rejected so header names stay unique.
//
// If leftField or rightField does not exist in its csv, or if the result
// would have duplicate headers, it returns a copy of s and false. The source
// csvs are not modified and the result shares no rows with them. When
// leftField and rightField have the same name, Join is equivalent to
// JoinByField.
func (s SimpleCsv) Join(other SimpleCsv, leftField, rightField string) (SimpleCsv, bool) {
	return s.joinAll(other, leftField, rightField, false, false)
}

// LeftJoin is like Join, but the rows of s without a match in other are also
// in the result, with empty strings in the columns of other.
func (s SimpleCsv) LeftJoin(other SimpleCsv, leftField, rightField string) (SimpleCsv, bool) {
	return s.joinAll(other, leftField, rightField, true, false)
}

// JoinE is like Join, but it returns an error describing the reason of the
// failure instead of false. On failure it returns an independent copy of s.
func (s SimpleCsv) JoinE(other SimpleCsv, leftField, rightField string) (SimpleCsv, error) {
	return s.joinOnE(other, []string{leftField}, []string{rightField}, false, false)
}

// LeftJoinE is like LeftJoin, but it returns an error describing the reason
// of the failure instead of false. On failure it returns an independent copy
// of s.
func (s SimpleCsv) LeftJoinE(other SimpleCsv, leftField, rightField string) (SimpleCsv, error) {
	return s.joinOnE(other, []string{leftField}, []string{rightField}, true, false)
}

// RightJoin returns a new csv with the right join of s and other by the
// column leftField of s and the column rightField of other, which may have
// different names. Every row of other is in the result: rows with a matching
// join value in s are joined to all their matches, and rows of other without
// a match are included with empty strings in the columns of s. The rows of s
// without a match in other are not in the result.
//
// Join values are compared exactly: case matters ("Ana" != "ana").
// Rows shorter than the join column position are treated as having an empty
// string in that cell, and cells beyond the header width are not in the
// result.
//
// The result headers are the headers of s followed by the headers of other
// excluding the join column rightField (the left join column is kept). Non-join
// columns of other must not share names with columns of s; a collision is
// rejected so header names stay unique.
//
// The result rows are in the order of s, with the rows of other that had no
// match appended at the end.
//
// If leftField or rightField does not exist in its csv, or if the result
// would have duplicate headers, it returns a copy of s and false. The source
// csvs are not modified and the result shares no rows with them. When
// leftField and rightField have the same name, RightJoin is equivalent to
// RightJoinByField.
func (s SimpleCsv) RightJoin(other SimpleCsv, leftField, rightField string) (SimpleCsv, bool) {
	return s.joinAll(other, leftField, rightField, false, true)
}

// FullJoin returns a new csv with the full (outer) join of s and other by
// the column leftField of s and the column rightField of other, which may
// have different names. Every row of s and every row of other is in the
// result: rows with a matching join value are joined to all their matches,
// and rows without a match are included with empty strings in the columns of
// the other side.
//
// Join values are compared exactly: case matters ("Ana" != "ana").
// Rows shorter than the join column position are treated as having an empty
// string in that cell, and cells beyond the header width are not in the
// result.
//
// The result headers are the headers of s followed by the headers of other
// excluding the join column rightField (the left join column is kept). Non-join
// columns of other must not share names with columns of s; a collision is
// rejected so header names stay unique.
//
// The result rows are in the order of s, with the rows of other that had no
// match appended at the end.
//
// If leftField or rightField does not exist in its csv, or if the result
// would have duplicate headers, it returns a copy of s and false. The source
// csvs are not modified and the result shares no rows with them. When
// leftField and rightField have the same name, FullJoin is equivalent to
// FullJoinByField.
func (s SimpleCsv) FullJoin(other SimpleCsv, leftField, rightField string) (SimpleCsv, bool) {
	return s.joinAll(other, leftField, rightField, true, true)
}

// RightJoinByField is like RightJoin, joining by the column fieldName, which
// must exist in both csvs.
func (s SimpleCsv) RightJoinByField(other SimpleCsv, fieldName string) (SimpleCsv, bool) {
	return s.joinAll(other, fieldName, fieldName, false, true)
}

// FullJoinByField is like FullJoin, joining by the column fieldName, which
// must exist in both csvs.
func (s SimpleCsv) FullJoinByField(other SimpleCsv, fieldName string) (SimpleCsv, bool) {
	return s.joinAll(other, fieldName, fieldName, true, true)
}

// JoinOn returns a new csv with the inner join of s and other by several
// columns at once. The i-th name of leftFields must be a column of s and the
// i-th name of rightFields a column of other: a row of s is joined to a row
// of other when the values of all its left key columns are exactly equal to
// the values of the corresponding right key columns. Only the rows whose
// join key exists in both csvs are in the result. If the same join key shows
// up more than once, all combinations of rows are in the result.
//
// Join values are compared exactly: case matters ("Ana" != "ana").
// Rows shorter than a join column position are treated as having an empty
// string in that cell, and cells beyond the header width are not in the
// result. The join key is built with a length-prefixed encoding, so values
// containing separators cannot collide.
//
// The result headers are the headers of s followed by the headers of other
// excluding all the right key columns (the left key columns are kept).
// Non-key columns of other must not share names with columns of s; a
// collision is rejected so header names stay unique.
//
// If s or other is empty, leftFields and rightFields have different lengths
// or are empty, one of the fields does not exist in its csv, or the result
// would have duplicate headers, it returns a copy of s and false. The source
// csvs are not modified and the result shares no rows with them. When both
// field lists have a single name, JoinOn is equivalent to Join.
func (s SimpleCsv) JoinOn(other SimpleCsv, leftFields, rightFields []string) (SimpleCsv, bool) {
	return s.joinOn(other, leftFields, rightFields, false, false)
}

// LeftJoinOn is like JoinOn, but the rows of s without a match in other are
// also in the result, with empty strings in the columns of other.
func (s SimpleCsv) LeftJoinOn(other SimpleCsv, leftFields, rightFields []string) (SimpleCsv, bool) {
	return s.joinOn(other, leftFields, rightFields, true, false)
}

// RightJoinOn is like JoinOn, but every row of other is in the result: rows
// with a matching join key in s are joined to all their matches, and rows of
// other without a match are included with empty strings in the columns of s.
// The rows of s without a match in other are not in the result. The rows of
// other that had no match are appended at the end.
func (s SimpleCsv) RightJoinOn(other SimpleCsv, leftFields, rightFields []string) (SimpleCsv, bool) {
	return s.joinOn(other, leftFields, rightFields, false, true)
}

// FullJoinOn is like JoinOn, but every row of s and every row of other is in
// the result: rows with a matching join key are joined to all their matches,
// and rows without a match are included with empty strings in the columns of
// the other side.
func (s SimpleCsv) FullJoinOn(other SimpleCsv, leftFields, rightFields []string) (SimpleCsv, bool) {
	return s.joinOn(other, leftFields, rightFields, true, true)
}

// joinAll joins s and other by the column leftField of s and the column
// rightField of other. If keepLeftUnmatched is true, the rows of s without a
// match in other are kept, with empty strings in the columns of other. If
// keepRightUnmatched is true, the rows of other without a match in s are
// kept, with empty strings in the columns of s, appended after all the rows
// of s.
func (s SimpleCsv) joinAll(other SimpleCsv, leftField, rightField string, keepLeftUnmatched, keepRightUnmatched bool) (SimpleCsv, bool) {
	return s.joinOn(other, []string{leftField}, []string{rightField}, keepLeftUnmatched, keepRightUnmatched)
}

// joinOn is the join used by JoinOn and, with single-name field lists, by
// all the other joins: it joins s and other by the columns leftFields of s
// and rightFields of other, matched positionally. The two lists must have
// the same length and at least one name, and every name must exist as a
// header in its csv. The join key of a row is built from the cells of the
// named columns with encodeKey, so values containing separators cannot
// collide. If keepLeftUnmatched is true, the rows of s without a match in
// other are kept, with empty strings in the columns of other. If
// keepRightUnmatched is true, the rows of other without a match in s are
// kept, with empty strings in the columns of s, appended after all the rows
// of s.
func (s SimpleCsv) joinOn(other SimpleCsv, leftFields, rightFields []string, keepLeftUnmatched, keepRightUnmatched bool) (SimpleCsv, bool) {
	result, err := s.joinOnE(other, leftFields, rightFields, keepLeftUnmatched, keepRightUnmatched)
	return result, err == nil
}

// joinOnE is like joinOn, but it returns an error describing the reason of
// the failure instead of false. On failure it returns an independent copy of
// s and a non-nil error whose message is prefixed with "simplecsv:".
func (s SimpleCsv) joinOnE(other SimpleCsv, leftFields, rightFields []string, keepLeftUnmatched, keepRightUnmatched bool) (SimpleCsv, error) {
	if len(s) == 0 {
		return s.failE(errors.New("simplecsv: cannot join: left csv is empty"))
	}
	if len(other) == 0 {
		return s.failE(errors.New("simplecsv: cannot join: right csv is empty"))
	}
	if len(leftFields) == 0 {
		return s.failE(errors.New("simplecsv: cannot join: at least one key field is required"))
	}
	if len(leftFields) != len(rightFields) {
		return s.failE(fmt.Errorf("simplecsv: cannot join: %d left key fields but %d right key fields", len(leftFields), len(rightFields)))
	}

	leftPositions := make([]int, len(leftFields))
	rightPositions := make([]int, len(rightFields))
	for i := range leftFields {
		leftPositions[i] = s.GetHeaderPosition(leftFields[i])
		if leftPositions[i] == -1 {
			return s.failE(fmt.Errorf("simplecsv: cannot join: key field %q not found in the left csv", leftFields[i]))
		}
		rightPositions[i] = other.GetHeaderPosition(rightFields[i])
		if rightPositions[i] == -1 {
			return s.failE(fmt.Errorf("simplecsv: cannot join: key field %q not found in the right csv", rightFields[i]))
		}
	}

	headers := joinHeaders(s[0], other[0], rightPositions)
	if colliding := collidingHeaders(s[0], other[0], rightPositions); len(colliding) > 0 {
		return s.failE(fmt.Errorf("simplecsv: cannot join: right columns %q collide with left columns", strings.Join(colliding, ", ")))
	}

	// Index the rows of other by join key once, so each row of s is
	// matched without scanning other again.
	otherIndex := make(map[string][]int)
	for i := 1; i < len(other); i++ {
		key := joinKey(other[i], rightPositions)
		otherIndex[key] = append(otherIndex[key], i)
	}

	joined := SimpleCsv{headers}
	matchedRight := make([]bool, len(other))
	for i := 1; i < len(s); i++ {
		matches := otherIndex[joinKey(s[i], leftPositions)]
		if len(matches) == 0 && keepLeftUnmatched {
			joined = append(joined, joinRow(s[i], nil, len(s[0]), len(other[0]), rightPositions))
		}
		for _, j := range matches {
			matchedRight[j] = true
			joined = append(joined, joinRow(s[i], other[j], len(s[0]), len(other[0]), rightPositions))
		}
	}
	if keepRightUnmatched {
		for j := 1; j < len(other); j++ {
			if !matchedRight[j] {
				joined = append(joined, joinRow(nil, other[j], len(s[0]), len(other[0]), rightPositions))
			}
		}
	}
	return joined, nil
}

// joinHeaders returns the headers of a join: the headers of left followed by
// the headers of right excluding the join columns at the positions in
// rightPositions.
func joinHeaders(left, right []string, rightPositions []int) []string {
	headers := make([]string, 0, len(left)+len(right)-len(rightPositions))
	headers = append(headers, left...)
	for k, header := range right {
		if !isJoinPosition(k, rightPositions) {
			headers = append(headers, header)
		}
	}
	return headers
}

// collidingHeaders returns the right headers, excluding the join columns at
// the positions in rightPositions, that also appear in left: joining would
// produce duplicate header names, breaking the unique-header invariant.
func collidingHeaders(left, right []string, rightPositions []int) []string {
	leftSet := make(map[string]struct{}, len(left))
	for _, header := range left {
		leftSet[header] = struct{}{}
	}
	var colliding []string
	for k, header := range right {
		if !isJoinPosition(k, rightPositions) {
			if _, exists := leftSet[header]; exists {
				colliding = append(colliding, header)
			}
		}
	}
	return colliding
}

// joinRow builds a joined row with the leftLen first cells of left and one
// cell per right column, excluding the join columns at the positions in
// rightPositions. Missing cells in ragged rows are empty strings. A nil
// right produces an all empty right side (left join without match).
func joinRow(left, right []string, leftLen, rightLen int, rightPositions []int) []string {
	row := make([]string, 0, leftLen+rightLen-len(rightPositions))
	for k := 0; k < leftLen; k++ {
		row = append(row, joinCell(left, k))
	}
	for k := 0; k < rightLen; k++ {
		if !isJoinPosition(k, rightPositions) {
			row = append(row, joinCell(right, k))
		}
	}
	return row
}

// joinKey builds the collision-free join key of a row from the cells at the
// named positions: every cell is prefixed with its byte length, so values
// containing separators cannot merge into the same key.
func joinKey(row []string, positions []int) string {
	parts := make([]string, len(positions))
	for i, position := range positions {
		parts[i] = joinCell(row, position)
	}
	return encodeKey(parts...)
}

// isJoinPosition reports whether position is one of the join column
// positions.
func isJoinPosition(position int, positions []int) bool {
	for _, p := range positions {
		if p == position {
			return true
		}
	}
	return false
}

// joinCell returns the cell at position in a row, or an empty string if the
// row is shorter than position (ragged row).
func joinCell(row []string, position int) string {
	if position < len(row) {
		return row[position]
	}
	return ""
}