github.com/matrixorigin/matrixone@v0.7.0/pkg/vm/engine/tae/dataio/blockio/writer.go (about)

     1  // Copyright 2021 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package blockio
    16  
    17  import (
    18  	"context"
    19  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    20  	"github.com/matrixorigin/matrixone/pkg/objectio"
    21  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers"
    22  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/index"
    23  )
    24  
    25  type Writer struct {
    26  	writer   objectio.Writer
    27  	fs       *objectio.ObjectFS
    28  	writeCxt context.Context
    29  }
    30  
    31  func NewWriter(ctx context.Context, fs *objectio.ObjectFS, name string) *Writer {
    32  	writer, err := objectio.NewObjectWriter(name, fs.Service)
    33  	if err != nil {
    34  		panic(any(err))
    35  	}
    36  	return &Writer{
    37  		fs:       fs,
    38  		writer:   writer,
    39  		writeCxt: ctx,
    40  	}
    41  }
    42  
    43  func (w *Writer) WriteBlock(columns *containers.Batch) (block objectio.BlockObject, err error) {
    44  	bat := batch.New(true, columns.Attrs)
    45  	bat.Vecs = containers.UnmarshalToMoVecs(columns.Vecs)
    46  	block, err = w.writer.Write(bat)
    47  	return
    48  }
    49  
    50  func (w *Writer) WriteBlockAndZoneMap(batch *batch.Batch, idxs []uint16) (objectio.BlockObject, error) {
    51  	block, err := w.writer.Write(batch)
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  	for _, idx := range idxs {
    56  		var zoneMap objectio.IndexData
    57  		vec := containers.NewVectorWithSharedMemory(batch.Vecs[idx], true)
    58  		zm := index.NewZoneMap(batch.Vecs[idx].Typ)
    59  		ctx := new(index.KeysCtx)
    60  		ctx.Keys = vec
    61  		ctx.Count = batch.Vecs[idx].Length()
    62  		defer ctx.Keys.Close()
    63  		err = zm.BatchUpdate(ctx)
    64  		if err != nil {
    65  			return nil, err
    66  		}
    67  		buf, err := zm.Marshal()
    68  		if err != nil {
    69  			return nil, err
    70  		}
    71  		zoneMap, err = objectio.NewZoneMap(idx, buf)
    72  		if err != nil {
    73  			return nil, err
    74  		}
    75  		w.writer.WriteIndex(block, zoneMap)
    76  	}
    77  	return block, nil
    78  }
    79  
    80  func (w *Writer) Sync() ([]objectio.BlockObject, error) {
    81  	blocks, err := w.writer.WriteEnd(w.writeCxt)
    82  	return blocks, err
    83  }
    84  
    85  func (w *Writer) WriteIndex(
    86  	block objectio.BlockObject,
    87  	index objectio.IndexData) (err error) {
    88  	return w.writer.WriteIndex(block, index)
    89  }
    90  
    91  func (w *Writer) GetWriter() objectio.Writer {
    92  	return w.writer
    93  }