c/src/main/java/org/apache/arrow/c/ArrowArrayStream.java

/*
 * 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.arrow.c;

import static org.apache.arrow.c.NativeUtil.NULL;
import static org.apache.arrow.util.Preconditions.checkNotNull;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import org.apache.arrow.c.jni.CDataJniException;
import org.apache.arrow.c.jni.JniWrapper;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.ReferenceManager;
import org.apache.arrow.memory.util.MemoryUtil;

/**
 * C Stream Interface ArrowArrayStream.
 *
 * <p>Represents a wrapper for the following C structure:
 *
 * <pre>
 * struct ArrowArrayStream {
 *   int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out);
 *   int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out);
 *   const char* (*get_last_error)(struct ArrowArrayStream*);
 *   void (*release)(struct ArrowArrayStream*);
 *   void* private_data;
 * };
 * </pre>
 */
public class ArrowArrayStream implements BaseStruct {
  private static final int SIZE_OF = 40;
  private static final int INDEX_RELEASE_CALLBACK = 24;

  private ArrowBuf data;

  /** Snapshot of the ArrowArrayStream raw data. */
  public static class Snapshot {
    public long get_schema;
    public long get_next;
    public long get_last_error;
    public long release;
    public long private_data;

    /** Initialize empty ArrowArray snapshot. */
    public Snapshot() {
      get_schema = NULL;
      get_next = NULL;
      get_last_error = NULL;
      release = NULL;
      private_data = NULL;
    }
  }

  /**
   * Create ArrowArrayStream from an existing memory address.
   *
   * <p>The resulting ArrowArrayStream does not own the memory.
   *
   * @param memoryAddress Memory address to wrap
   * @return A new ArrowArrayStream instance
   */
  public static ArrowArrayStream wrap(long memoryAddress) {
    return new ArrowArrayStream(
        new ArrowBuf(ReferenceManager.NO_OP, null, ArrowArrayStream.SIZE_OF, memoryAddress));
  }

  /**
   * Create ArrowArrayStream by allocating memory.
   *
   * <p>The resulting ArrowArrayStream owns the memory.
   *
   * @param allocator Allocator for memory allocations
   * @return A new ArrowArrayStream instance
   */
  public static ArrowArrayStream allocateNew(BufferAllocator allocator) {
    ArrowArrayStream array = new ArrowArrayStream(allocator.buffer(ArrowArrayStream.SIZE_OF));
    array.markReleased();
    return array;
  }

  ArrowArrayStream(ArrowBuf data) {
    checkNotNull(data, "ArrowArrayStream initialized with a null buffer");
    this.data = data;
  }

  /** Mark the array as released. */
  public void markReleased() {
    directBuffer().putLong(INDEX_RELEASE_CALLBACK, NULL);
  }

  @Override
  public long memoryAddress() {
    checkNotNull(data, "ArrowArrayStream is already closed");
    return data.memoryAddress();
  }

  @Override
  public void release() {
    long address = memoryAddress();
    JniWrapper.get().releaseArrayStream(address);
  }

  /**
   * Get the schema of the stream.
   *
   * @param schema The ArrowSchema struct to output to
   * @throws IOException if the stream returns an error
   */
  public void getSchema(ArrowSchema schema) throws IOException {
    long address = memoryAddress();
    try {
      JniWrapper.get().getSchemaArrayStream(address, schema.memoryAddress());
    } catch (CDataJniException e) {
      throw new IOException("[errno " + e.getErrno() + "] " + e.getMessage());
    }
  }

  /**
   * Get the next batch in the stream.
   *
   * @param array The ArrowArray struct to output to
   * @throws IOException if the stream returns an error
   */
  public void getNext(ArrowArray array) throws IOException {
    long address = memoryAddress();
    try {
      JniWrapper.get().getNextArrayStream(address, array.memoryAddress());
    } catch (CDataJniException e) {
      throw new IOException("[errno " + e.getErrno() + "] " + e.getMessage());
    }
  }

  @Override
  public void close() {
    if (data != null) {
      data.close();
      data = null;
    }
  }

  private ByteBuffer directBuffer() {
    return MemoryUtil.directBuffer(memoryAddress(), ArrowArrayStream.SIZE_OF)
        .order(ByteOrder.nativeOrder());
  }

  /**
   * Take a snapshot of the ArrowArrayStream raw values.
   *
   * @return snapshot
   */
  public ArrowArrayStream.Snapshot snapshot() {
    ByteBuffer data = directBuffer();
    ArrowArrayStream.Snapshot snapshot = new ArrowArrayStream.Snapshot();
    snapshot.get_schema = data.getLong();
    snapshot.get_next = data.getLong();
    snapshot.get_last_error = data.getLong();
    snapshot.release = data.getLong();
    snapshot.private_data = data.getLong();
    return snapshot;
  }

  /** Write values from Snapshot to the underlying ArrowArrayStream memory buffer. */
  public void save(ArrowArrayStream.Snapshot snapshot) {
    directBuffer()
        .putLong(snapshot.get_schema)
        .putLong(snapshot.get_next)
        .putLong(snapshot.get_last_error)
        .putLong(snapshot.release)
        .putLong(snapshot.private_data);
  }
}