show result for v2 protocol (#21)

* show result for TransactionAsyncResult * enhance v2 show * enhance v2 show result * Update rai/show.go Co-authored-by: Nathan Daly <[email protected]> * Update rai/show.go Co-authored-by: Nathan Daly <[email protected]> * update * make arrowRelation unique by relation id * fix arrow relation test * Simplification to readArrowFiles to construct the columns in one pass (#22) heh, sorry i have one more suggestion. I guess you can do this in a single-pass. (Sorry, github wouldn't let me leave a code suggestion over this large section of the code) * adding ShowIO and v2 show test Co-authored-by: Nathan Daly <[email protected]>
RelationalAI · Aug 6, 2022 · a3f7263 · a3f7263
1 parent 6ac387b
commit a3f7263
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 8 deletions.
diff --git a/rai/client.go b/rai/client.go
@@ -349,12 +349,14 @@ func readArrowFiles(files []TransactionAsyncFile) ([]ArrowRelation, error) {
 			defer reader.Release()
 			for reader.Next() {
 				rec := reader.Record()
+				var columns [][]interface{}
 				for i := 0; i < int(rec.NumCols()); i++ {
 					data, _ := rec.Column(i).MarshalJSON()
-					var values []interface{}
-					json.Unmarshal(data, &values)
-					out = append(out, ArrowRelation{file.Name, values})
+					var column []interface{}
+					json.Unmarshal(data, &column)
+					columns = append(columns, column)
 				}
+				out = append(out, ArrowRelation{file.Name, columns})
 
 				rec.Retain()
 			}

diff --git a/rai/client_test.go b/rai/client_test.go
@@ -15,6 +15,7 @@
 package rai
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"net/http"
@@ -326,10 +327,12 @@ func TestExecuteAsync(t *testing.T) {
 	assert.Nil(t, err)
 
 	expectedResults := []ArrowRelation{
-		ArrowRelation{"/:output/Int64/Int64/Int64/Int64", []interface{}{1., 2., 3., 4., 5.}},
-		ArrowRelation{"/:output/Int64/Int64/Int64/Int64", []interface{}{1., 4., 9., 16., 25.}},
-		ArrowRelation{"/:output/Int64/Int64/Int64/Int64", []interface{}{1., 8., 27., 64., 125.}},
-		ArrowRelation{"/:output/Int64/Int64/Int64/Int64", []interface{}{1., 16., 81., 256., 625.}},
+		ArrowRelation{"/:output/Int64/Int64/Int64/Int64", [][]interface{}{
+			{1., 2., 3., 4., 5.},
+			{1., 4., 9., 16., 25.},
+			{1., 8., 27., 64., 125.},
+			{1., 16., 81., 256., 625.},
+		}},
 	}
 
 	assert.Equal(t, rsp.Results[0].Table, expectedResults[0].Table)
@@ -343,6 +346,13 @@ func TestExecuteAsync(t *testing.T) {
 	expectedProblems := []interface{}{}
 
 	assert.Equal(t, rsp.Problems, expectedProblems)
+
+	// also testing Show v2 result format
+	var io bytes.Buffer
+	rsp.ShowIO(&io)
+	expectedOutput := "/:output/Int64/Int64/Int64/Int64\n1, 1, 1, 1\n2, 4, 8, 16\n3, 9, 27, 81\n4, 16, 64, 256\n5, 25, 125, 625\n\n"
+
+	assert.Equal(t, io.String(), expectedOutput)
 }
 
 func findRelation(relations []Relation, colName string) *Relation {

diff --git a/rai/models.go b/rai/models.go
@@ -258,7 +258,7 @@ type TransactionAsyncFile struct {
 
 type ArrowRelation struct {
 	RelationID string
-	Table      []interface{}
+	Table      [][]interface{}
 }
 
 type TransactionAsyncSingleResponse struct {

diff --git a/rai/show.go b/rai/show.go
@@ -127,3 +127,42 @@ func (tx *TransactionResult) Show() {
 		}
 	}
 }
+
+func zip(lists ...[]interface{}) func() []interface{} {
+	zip := make([]interface{}, len(lists))
+	i := 0
+	return func() []interface{} {
+		for j := range lists {
+			if i >= len(lists[j]) {
+				return nil
+			}
+			zip[j] = lists[j][i]
+		}
+		i++
+		return zip
+	}
+}
+
+func (tx *TransactionAsyncResult) ShowIO(io io.Writer) {
+	for _, r := range tx.Results {
+		k := r.RelationID
+		v := r.Table
+		fmt.Fprintf(io, "%s\n", k)
+		iter := zip(v...)
+		for tuple := iter(); tuple != nil; tuple = iter() {
+			for i, element := range tuple {
+				if i > 0 {
+					fmt.Fprint(io, ", ")
+				}
+
+				fmt.Fprintf(io, "%v", element)
+			}
+			fmt.Fprintln(io)
+		}
+		fmt.Fprintln(io)
+	}
+}
+
+func (tx *TransactionAsyncResult) Show() {
+	tx.ShowIO(os.Stdout)
+}