org.apache.iceberg.ksyun.ks3.Ks3AppendOutputStream Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of iceberg-ksyun Show documentation
apache iceberg support ks3
The newest version!
/*
 * 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.
 */

package org.apache.iceberg.ksyun.ks3;

import com.ksyun.ks3.dto.ObjectMetadata;
import com.ksyun.ks3.service.Ks3;
import com.ksyun.ks3.service.request.AppendObjectRequest;
import com.ksyun.ks3.service.request.PutObjectRequest;
import org.apache.iceberg.io.PositionOutputStream;

import java.io.ByteArrayInputStream;
import java.nio.ByteBuffer;

/**
 * Use KS3 append API to write data.
 */
class Ks3AppendOutputStream extends PositionOutputStream {

  private final Ks3 client;

  private final Ks3URI uri;

  /**
   * Local bytes cache that avoid too many requests
   * 
   * Use {@link ByteBuffer} to maintain offset.
   */
  private final ByteBuffer localCache;

  /**
   * Pos for {@link PositionOutputStream}
   */
  private long pos;

  private long objectPos;


  private Ks3AppendOutputStream(Ks3 client, Ks3URI uri, byte[] localCache) {
    this.client = client;
    this.uri = uri;
    this.localCache = ByteBuffer.wrap(localCache);
  }

  /**
   * Use built-in 1 KiB byte buffer
   */
  static Ks3AppendOutputStream create(Ks3 client, Ks3URI uri) {
    return createWithBufferSize(client, uri, 1024);
  }

  /**
   * Create {@link PositionOutputStream} with specific buffer size.
   */
  static Ks3AppendOutputStream createWithBufferSize(Ks3 client, Ks3URI uri, int size) {
    return new Ks3AppendOutputStream(client, uri, new byte[size]);
  }

  /**
   * Write a byte. If buffer is full, upload the buffer.
   */
  @Override
  public void write(int b) {
    if (!checkBuffer(1)) {
      flush();
    }

    localCache.put((byte) b);
    pos += 1;
  }

  /**
   * Write a byte.
   * If buffer is full, upload the buffer.
   * If buffer size < input bytes, upload input bytes.
   */
  @Override
  public void write(byte[] b, int off, int len) {
    if (!checkBuffer(len)) {
      flush();
    }

    if (checkBuffer(len)) {
      localCache.put(b, off, len);
    } else {
      // if content > cache, directly flush itself.
      flushBuffer(b, off, len);
    }

    pos += len;
  }

  private boolean checkBuffer(int nextWrite) {
    return localCache.remaining() >= nextWrite;
  }

  private void flushBuffer(byte[] buffer, int offset, int length) {
    AppendObjectRequest appendObjectRequest = new AppendObjectRequest(uri.bucket(), uri.name(),
        new ByteArrayInputStream(buffer, offset, length), null);
    appendObjectRequest.setPosition(objectPos);
    client.appendObject(appendObjectRequest);
    objectPos += length;
  }

  /**
   * Pos of the file
   */
  @Override
  public long getPos() {
    return pos;
  }

  /**
   * Write cached bytes if present.
   */
  @Override
  public void flush() {
    if (localCache.remaining() < localCache.capacity()) {
      localCache.flip();
      flushBuffer(localCache.array(), localCache.arrayOffset(), localCache.remaining());
      localCache.clear();
    }
  }

  /**
   * Trigger flush() when closing stream.
   */
  @Override
  public void close() {
    flush();
  }
}