Allow IntSequenceTable to take a non-constant argument.

Author: nicktobey

memory/lookup_squence_table.go

@@ -18,6 +18,7 @@ var _ sql.TableNode = LookupSequenceTable{}
 // LookupSequenceTable is a variation of IntSequenceTable that supports lookups and implements sql.TableNode
 type LookupSequenceTable struct {
 	IntSequenceTable
+	length int64
 }
 
 func (s LookupSequenceTable) UnderlyingTable() sql.Table {
@@ -29,7 +30,15 @@ func (s LookupSequenceTable) NewInstance(ctx *sql.Context, db sql.Database, args
 	if err != nil {
 		return nil, err
 	}
-	return LookupSequenceTable{newIntSequenceTable.(IntSequenceTable)}, nil
+	lenExp, ok := args[1].(*expression.Literal)
+	if !ok {
+		return nil, fmt.Errorf("sequence table expects arguments to be literal expressions")
+	}
+	length, _, err := types.Int64.Convert(ctx, lenExp.Value())
+	if err != nil {
+		return nil, err
+	}
+	return LookupSequenceTable{newIntSequenceTable.(IntSequenceTable), length.(int64)}, nil
 }
 
 func (s LookupSequenceTable) String() string {
@@ -82,21 +91,21 @@ func (s LookupSequenceTable) Description() string {
 
 // Partitions is a sql.Table interface function that returns a partition of the data. This data has a single partition.
 func (s LookupSequenceTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return sql.PartitionsToPartitionIter(&sequencePartition{min: 0, max: int64(s.Len) - 1}), nil
+	return sql.PartitionsToPartitionIter(&sequencePartition{min: 0, max: s.length - 1}), nil
 }
 
 // PartitionRows is a sql.Table interface function that takes a partition and returns all rows in that partition.
 // This table has a partition for just schema changes, one for just data changes, and one for both.
 func (s LookupSequenceTable) PartitionRows(ctx *sql.Context, partition sql.Partition) (sql.RowIter, error) {
 	sp, ok := partition.(*sequencePartition)
 	if !ok {
-		return &SequenceTableFnRowIter{i: 0, n: s.Len}, nil
+		return &SequenceTableFnRowIter{i: 0, n: s.length}, nil
 	}
 	min := int64(0)
 	if sp.min > min {
 		min = sp.min
 	}
-	max := int64(s.Len) - 1
+	max := int64(s.length) - 1
 	if sp.max < max {
 		max = sp.max
 	}

memory/sequence_table.go

@@ -19,7 +19,7 @@ var _ sql.ExecSourceRel = IntSequenceTable{}
 type IntSequenceTable struct {
 	db   sql.Database
 	name string
-	Len  int64
+	Len  sql.Expression
 }
 
 func (s IntSequenceTable) NewInstance(ctx *sql.Context, db sql.Database, args []sql.Expression) (sql.Node, error) {
@@ -34,15 +34,11 @@ func (s IntSequenceTable) NewInstance(ctx *sql.Context, db sql.Database, args []
 	if !ok {
 		return nil, fmt.Errorf("sequence table expects 1st argument to be column name")
 	}
-	lenExp, ok := args[1].(*expression.Literal)
-	if !ok {
-		return nil, fmt.Errorf("sequence table expects arguments to be literal expressions")
+	lenExp := args[1]
+	if !sql.IsNumberType(lenExp.Type()) {
+		return nil, fmt.Errorf("sequence table expects length argument to be a number")
 	}
-	length, _, err := types.Int64.Convert(ctx, lenExp.Value())
-	if !ok {
-		return nil, fmt.Errorf("%w; sequence table expects 2nd argument to be a sequence length integer", err)
-	}
-	return IntSequenceTable{db: db, name: name, Len: length.(int64)}, nil
+	return IntSequenceTable{db: db, name: name, Len: lenExp}, nil
 }
 
 func (s IntSequenceTable) Resolved() bool {
@@ -85,8 +81,20 @@ func (s IntSequenceTable) Children() []sql.Node {
 	return []sql.Node{}
 }
 
-func (s IntSequenceTable) RowIter(_ *sql.Context, _ sql.Row) (sql.RowIter, error) {
-	rowIter := &SequenceTableFnRowIter{i: 0, n: s.Len}
+func (s IntSequenceTable) RowIter(ctx *sql.Context, row sql.Row) (sql.RowIter, error) {
+	iterLen, err := s.Len.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+	iterLenVal, ok, err := types.Int64.Convert(ctx, iterLen)
+	if err != nil {
+		return nil, err
+	}
+	if !ok {
+		return nil, fmt.Errorf("sequence table expects integer argument")
+	}
+
+	rowIter := &SequenceTableFnRowIter{i: 0, n: iterLenVal.(int64)}
 	return rowIter, nil
 }