github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/partitioner.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package colexec
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    21  )
    22  
    23  // NewWindowSortingPartitioner creates a new colexec.Operator that orders input
    24  // first based on the partitionIdxs columns and second on ordCols (i.e. it
    25  // handles both PARTITION BY and ORDER BY clauses of a window function) and
    26  // puts true in partitionColIdx'th column (which is appended if needed) for
    27  // every tuple that is the first within its partition.
    28  func NewWindowSortingPartitioner(
    29  	allocator *colmem.Allocator,
    30  	input colexecbase.Operator,
    31  	inputTyps []*types.T,
    32  	partitionIdxs []uint32,
    33  	ordCols []execinfrapb.Ordering_Column,
    34  	partitionColIdx int,
    35  	createDiskBackedSorter func(input colexecbase.Operator, inputTypes []*types.T, orderingCols []execinfrapb.Ordering_Column) (colexecbase.Operator, error),
    36  ) (op colexecbase.Operator, err error) {
    37  	partitionAndOrderingCols := make([]execinfrapb.Ordering_Column, len(partitionIdxs)+len(ordCols))
    38  	for i, idx := range partitionIdxs {
    39  		partitionAndOrderingCols[i] = execinfrapb.Ordering_Column{ColIdx: idx}
    40  	}
    41  	copy(partitionAndOrderingCols[len(partitionIdxs):], ordCols)
    42  	input, err = createDiskBackedSorter(input, inputTyps, partitionAndOrderingCols)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  
    47  	var distinctCol []bool
    48  	input, distinctCol, err = OrderedDistinctColsToOperators(input, partitionIdxs, inputTyps)
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  
    53  	input = newVectorTypeEnforcer(allocator, input, types.Bool, partitionColIdx)
    54  	return &windowSortingPartitioner{
    55  		OneInputNode:    NewOneInputNode(input),
    56  		allocator:       allocator,
    57  		distinctCol:     distinctCol,
    58  		partitionColIdx: partitionColIdx,
    59  	}, nil
    60  }
    61  
    62  type windowSortingPartitioner struct {
    63  	OneInputNode
    64  
    65  	allocator *colmem.Allocator
    66  	// distinctCol is the output column of the chain of ordered distinct
    67  	// operators in which true will indicate that a new partition begins with the
    68  	// corresponding tuple.
    69  	distinctCol     []bool
    70  	partitionColIdx int
    71  }
    72  
    73  func (p *windowSortingPartitioner) Init() {
    74  	p.input.Init()
    75  }
    76  
    77  func (p *windowSortingPartitioner) Next(ctx context.Context) coldata.Batch {
    78  	b := p.input.Next(ctx)
    79  	if b.Length() == 0 {
    80  		return coldata.ZeroBatch
    81  	}
    82  	partitionVec := b.ColVec(p.partitionColIdx)
    83  	if partitionVec.MaybeHasNulls() {
    84  		// We need to make sure that there are no left over null values in the
    85  		// output vector.
    86  		partitionVec.Nulls().UnsetNulls()
    87  	}
    88  	partitionCol := partitionVec.Bool()
    89  	sel := b.Selection()
    90  	if sel != nil {
    91  		for i := 0; i < b.Length(); i++ {
    92  			partitionCol[sel[i]] = p.distinctCol[sel[i]]
    93  		}
    94  	} else {
    95  		copy(partitionCol, p.distinctCol[:b.Length()])
    96  	}
    97  	return b
    98  }