github.com/apache/arrow/go/v14@v14.0.1/parquet/internal/utils/write_utils.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one
     2  // or more contributor license agreements.  See the NOTICE file
     3  // distributed with this work for additional information
     4  // regarding copyright ownership.  The ASF licenses this file
     5  // to you under the Apache License, Version 2.0 (the
     6  // "License"); you may not use this file except in compliance
     7  // with the License.  You may obtain a copy of the License at
     8  //
     9  // http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  
    17  package utils
    18  
    19  import "io"
    20  
    21  // WriterTell is an interface that adds a Tell function to an io.Writer
    22  type WriterTell interface {
    23  	io.Writer
    24  	Tell() int64
    25  }
    26  
    27  // WriteCloserTell is an interface adding a Tell function to a WriteCloser
    28  // so if the underlying writer has a Close function, it is exposed and not
    29  // hidden.
    30  type WriteCloserTell interface {
    31  	io.WriteCloser
    32  	Tell() int64
    33  }
    34  
    35  // TellWrapper wraps any io.Writer to add a Tell function that tracks
    36  // the position based on calls to Write. It does not take into account
    37  // any calls to Seek or any Writes that don't go through the TellWrapper
    38  type TellWrapper struct {
    39  	io.Writer
    40  	pos int64
    41  }
    42  
    43  // Close makes TellWrapper an io.Closer so that calling Close
    44  // will also call Close on the wrapped writer if it has a Close function.
    45  func (w *TellWrapper) Close() error {
    46  	if closer, ok := w.Writer.(io.WriteCloser); ok {
    47  		return closer.Close()
    48  	}
    49  	return nil
    50  }
    51  
    52  func (w *TellWrapper) Tell() int64 { return w.pos }
    53  func (w *TellWrapper) Write(p []byte) (n int, err error) {
    54  	n, err = w.Writer.Write(p)
    55  	w.pos += int64(n)
    56  	return
    57  }