Add NumStripes() and SelectStripe() to Reader and Cursor (#56)

* Add NumStripes() and GetStripe() to Reader and Cursor

* Fix stripe numbers to start from idx 0 instead of 1

* Address PR comments

* Remove error variable declaration

* Remove go-cmp from cursor test

Author: athum

cursor.go

@@ -44,6 +44,17 @@ func (c *Cursor) Select(fields ...string) *Cursor {
 	return c
 }
 
+// SelectStripe retrieves the stream information for the specified stripe.
+func (c *Cursor) SelectStripe(n int) error {
+	stripe, err := c.Reader.getStripe(n, c.included...)
+	if err != nil {
+		return err
+	}
+	c.Stripe = stripe
+	c.stripeOffset = n
+	return c.prepareStreamReaders()
+}
+
 // prepareStreamReaders prepares TreeReaders for each of the columns
 // that will be read.
 func (c *Cursor) prepareStreamReaders() error {
@@ -65,10 +76,11 @@ func (c *Cursor) prepareNextStripe() error {
 	// and creating the required readers for each of the
 	// required columns.
 	var err error
-	c.Stripe, err = c.Reader.getStripe(c.stripeOffset, c.included...)
+	stripe, err := c.Reader.getStripe(c.stripeOffset, c.included...)
 	if err != nil {
 		return err
 	}
+	c.Stripe = stripe
 	c.stripeOffset++ // Increment in order to fetch the next stripe.
 	return c.prepareStreamReaders()
 }

cursor_test.go

@@ -93,7 +93,6 @@ func TestCursorResets(t *testing.T) {
 
 }
 
-
 func TestCursorSelectError(t *testing.T) {
 
 	r, err := Open("./examples/demo-11-zlib.orc")
@@ -104,23 +103,75 @@ func TestCursorSelectError(t *testing.T) {
 
 	// Try to select a column that doesn't exist.
 	c := r.Select("notfound")
-	
+
 	var hasNext bool
 	for c.Next() {
 		hasNext = true
 	}
-	
+
 	if hasNext {
 		t.Errorf("Next returned true, expected false")
 	}
-	
+
 	err = c.Err()
 	if err == nil {
 		t.Errorf("Expected error")
 	}
 	if err.Error() != "no field with name: notfound" {
 		t.Errorf("Unexpected error: %s", err.Error())
 	}
-	
-	
+
+}
+
+func TestCursorSelectStripe(t *testing.T) {
+	r, err := Open("./examples/demo-11-none.orc")
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+
+	expectedStripes := 385
+
+	numStripes, err := r.NumStripes()
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+
+	if numStripes != expectedStripes {
+		t.Fatalf("Expected %d stripes, got %d", expectedStripes, numStripes)
+	}
+
+	cols := r.Schema().Columns()
+	c := r.Select(cols...)
+
+	err = c.SelectStripe(numStripes - 1)
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+
+	var row []interface{}
+	for c.Next() {
+		row = c.Row()
+	}
+
+	expectedLastRow := []interface{}{
+		int64(1920800),
+		"F",
+		"U",
+		"Unknown",
+		int64(10000),
+		"Unknown",
+		int64(6),
+		int64(6),
+		int64(6),
+	}
+
+	if len(row) != len(expectedLastRow) {
+		t.Fatalf("Expected %d elements, got %d", len(expectedLastRow), len(row))
+	}
+
+	for idx, e := range expectedLastRow {
+		if e != row[idx] {
+			t.Fatalf("Expected %v, got %v", e, row[idx])
+		}
+	}
 }

reader.go

@@ -341,6 +341,15 @@ func (r *Reader) NumRows() int {
 	return int(r.footer.GetNumberOfRows())
 }
 
+func (r *Reader) NumStripes() (int, error) {
+	stripes, err := r.getStripes()
+	if err != nil {
+		return 0, err
+	}
+
+	return len(stripes), nil
+}
+
 type Stripe struct {
 	included []int
 	*proto.StripeInformation

reader_test.go

@@ -37,3 +37,23 @@ func TestReadNullAtEnd(t *testing.T) {
 		}
 	}
 }
+
+func TestNumStripes(t *testing.T) {
+	expectedStripes := 385
+
+	r, err := Open("examples/demo-11-none.orc")
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+
+	defer r.Close()
+
+	n, err := r.NumStripes()
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+
+	if n != expectedStripes {
+		t.Fatalf("Expected %d stripes, got %d", expectedStripes, n)
+	}
+}