Drive coordinator upload and reduction from the workload catalog

coordinator/internal/reducer/concat.go

@@ -0,0 +1,57 @@
+// Package reducer contains deterministic, coordinator-side result reductions.
+package reducer
+
+import (
+	"bytes"
+	"encoding/csv"
+	"errors"
+	"fmt"
+	"io"
+)
+
+// ReduceOrderedConcat concatenates worker partial tables in shard order into a
+// single table with one header. Every partial must carry the same header as the
+// first partial and rows of the same width; anything else fails the job closed.
+func ReduceOrderedConcat(partials []io.Reader) ([]byte, error) {
+	var out bytes.Buffer
+	writer := csv.NewWriter(&out)
+	var firstHeader []string
+	for _, partial := range partials {
+		reader := csv.NewReader(partial)
+		header, err := reader.Read()
+		if err != nil {
+			return nil, fmt.Errorf("read partial header: %w", err)
+		}
+		if len(header) == 0 {
+			return nil, fmt.Errorf("partial result has an empty header")
+		}
+		if firstHeader == nil {
+			firstHeader = header
+			if err := writer.Write(header); err != nil {
+				return nil, err
+			}
+		} else if !equalStrings(header, firstHeader) {
+			return nil, fmt.Errorf("partial result has an inconsistent header")
+		}
+		for {
+			row, err := reader.Read()
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			if err != nil {
+				return nil, fmt.Errorf("read partial row: %w", err)
+			}
+			if len(row) != len(header) {
+				return nil, fmt.Errorf("partial result has a row with an inconsistent width")
+			}
+			if err := writer.Write(row); err != nil {
+				return nil, err
+			}
+		}
+	}
+	writer.Flush()
+	if err := writer.Error(); err != nil {
+		return nil, err
+	}
+	return out.Bytes(), nil
+}

coordinator/internal/reducer/concat_test.go

@@ -0,0 +1,64 @@
+package reducer
+
+import (
+	"io"
+	"strings"
+	"testing"
+)
+
+func TestReduceOrderedConcatJoinsPartialsInOrderWithOneHeader(t *testing.T) {
+	first := strings.NewReader("chembl_id,canonical_smiles\nA,CC\nB,CCC\n")
+	second := strings.NewReader("chembl_id,canonical_smiles\nC,CCCC\n")
+
+	output, err := ReduceOrderedConcat([]io.Reader{first, second})
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := "chembl_id,canonical_smiles\nA,CC\nB,CCC\nC,CCCC\n"
+	if string(output) != want {
+		t.Fatalf("output = %q, want %q", output, want)
+	}
+}
+
+func TestReduceOrderedConcatIsDeterministicAcrossInputOrder(t *testing.T) {
+	left := strings.NewReader("id,rows\nA,1\nB,2\n")
+	right := strings.NewReader("id,rows\nC,3\n")
+
+	first, err := ReduceOrderedConcat([]io.Reader{left, right})
+	if err != nil {
+		t.Fatal(err)
+	}
+	left, right = strings.NewReader("id,rows\nA,1\nB,2\n"), strings.NewReader("id,rows\nC,3\n")
+	second, err := ReduceOrderedConcat([]io.Reader{left, right})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(first) != string(second) {
+		t.Fatalf("concat is not deterministic: %q != %q", first, second)
+	}
+}
+
+func TestReduceOrderedConcatRejectsInconsistentHeaders(t *testing.T) {
+	first := strings.NewReader("a,b\n1,2\n")
+	second := strings.NewReader("a,c\n1,2\n")
+	if _, err := ReduceOrderedConcat([]io.Reader{first, second}); err == nil {
+		t.Fatal("inconsistent headers must fail")
+	}
+}
+
+func TestReduceOrderedConcatRejectsRaggedRows(t *testing.T) {
+	partial := strings.NewReader("a,b\n1,2,3\n")
+	if _, err := ReduceOrderedConcat([]io.Reader{partial}); err == nil {
+		t.Fatal("ragged rows must fail")
+	}
+}
+
+func TestReduceOrderedConcatEmptyPartials(t *testing.T) {
+	output, err := ReduceOrderedConcat(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(output) != 0 {
+		t.Fatalf("empty input must produce empty output, got %q", output)
+	}
+}