avro_stream_internal.h

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#pragma once
 
#include <arrow/io/interfaces.h>
#include <avro/Stream.hh>
 
namespace iceberg::avro {
 
class AvroInputStream : public ::avro::SeekableInputStream {
 public:
  explicit AvroInputStream(std::shared_ptr<::arrow::io::RandomAccessFile> input_stream,
                           int64_t buffer_size);
 
  ~AvroInputStream() override;
 
  bool next(const uint8_t** data, size_t* len) override;
 
  void backup(size_t len) override;
 
  void skip(size_t len) override;
 
  size_t byteCount() const override;
 
  void seek(int64_t position) override;
 
 private:
  std::shared_ptr<::arrow::io::RandomAccessFile> input_stream_;
  const int64_t buffer_size_;
  std::vector<uint8_t> buffer_;
  size_t byte_count_ = 0;       // bytes read from the input stream
  size_t buffer_pos_ = 0;       // next position to read in the buffer
  size_t available_bytes_ = 0;  // bytes available in the buffer
};
 
class AvroOutputStream : public ::avro::OutputStream {
 public:
  explicit AvroOutputStream(std::shared_ptr<::arrow::io::OutputStream> output_stream,
                            int64_t buffer_size);
 
  ~AvroOutputStream() override;
 
  bool next(uint8_t** data, size_t* len) override;
 
  void backup(size_t len) override;
 
  uint64_t byteCount() const override;
 
  void flush() override;
 
  const std::shared_ptr<::arrow::io::OutputStream>& arrow_output_stream() const;
 
 private:
  std::shared_ptr<::arrow::io::OutputStream> output_stream_;
  const int64_t buffer_size_;
  std::vector<uint8_t> buffer_;
  size_t buffer_pos_ = 0;       // position in the buffer
  uint64_t flushed_bytes_ = 0;  // bytes flushed to the output stream
};
 
}  // namespace iceberg::avro