Arrow Flight and Deephaven

import io.deephaven.extensions.barrage.chunk.DefaultChunkWriterFactory

DefaultChunkWriterFactory.INSTANCE.disableToStringUnknownTypes()

import io.deephaven.engine.table.Table
import io.deephaven.extensions.barrage.util.BarrageUtil
import org.apache.arrow.vector.types.TimeUnit
import org.apache.arrow.vector.types.Types
import org.apache.arrow.vector.types.pojo.ArrowType
import org.apache.arrow.vector.types.pojo.Field
import org.apache.arrow.vector.types.pojo.FieldType
import org.apache.arrow.vector.types.pojo.Schema

Table t = emptyTable(32).update("D = ii * 1_000_000_000L")

Schema curr_schema = BarrageUtil.schemaFromTable(t)
Field curr_field = curr_schema.getFields().get(0)
FieldType curr_ftype = curr_field.getFieldType()
Map<String, String> curr_metadata = curr_field.getMetadata()
// note that we could replace the destination metadata deephaven type, too, if the remote client is a Deephaven worker

ArrowType new_type = new ArrowType.Duration(TimeUnit.MILLISECOND)
FieldType new_ftype = new FieldType(curr_ftype.isNullable(), new_type, curr_ftype.getDictionary(), curr_metadata);
List<Field> new_fields = List.of(new Field(curr_field.getName(), new_ftype, curr_field.getChildren()))
Schema new_schema = new Schema(new_fields, curr_schema.getCustomMetadata())

// now set the table attribute to advertise how to write this on the wire
t = t.withAttributes(Map.of(Table.BARRAGE_SCHEMA_ATTRIBUTE, new_schema))

import java.time.Duration
import io.deephaven.util.QueryConstants
import io.deephaven.chunk.*
import io.deephaven.chunk.attributes.*
import io.deephaven.extensions.barrage.chunk.*
import org.apache.arrow.vector.types.TimeUnit
import org.apache.arrow.vector.types.pojo.ArrowType


long factorForTimeUnit(TimeUnit unit) {
    switch (unit) {
        case TimeUnit.NANOSECOND:
            return 1L
        case TimeUnit.MICROSECOND:
            return 1000L
        case TimeUnit.MILLISECOND:
            return 1000 * 1000L
        case TimeUnit.SECOND:
            return 1000 * 1000 * 1000L
        default:
            throw new RuntimeException("unknown time unit")
    }
}

DefaultChunkReaderFactory.INSTANCE.register(
        ArrowType.ArrowTypeID.Duration,
        long.class,
        (arrowType, typeInfo, options) -> {
            final long factor = factorForTimeUnit(((ArrowType.Duration) arrowType).getUnit());
            return new LongChunkReader(options) {
                @Override def WritableLongChunk<Values> readChunk(
                        Iterator<ChunkWriter.FieldNodeInfo> fieldNodeIter,
                        PrimitiveIterator.OfLong bufferInfoIter,
                        DataInput is,
                        WritableChunk<Values> outChunk,
                        final int outOffset,
                        final int totalRows) throws IOException {
                    final WritableLongChunk<Values> values =
                            super.readChunk(fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
                    for (int ii = outOffset; ii < values.size(); ++ii) {
                        if (!values.isNull(ii)) {
                            values.set(ii, values.get(ii) * factor);
                        }
                    }
                    return values;
                }
            };
        }
)

DefaultChunkWriterFactory.INSTANCE.register(
        ArrowType.ArrowTypeID.Duration,
        long.class,
        typeInfo -> {
            final long factor = factorForTimeUnit(((ArrowType.Duration) typeInfo.arrowField().getType()).getUnit());
            return new LongChunkWriter<>((source) -> {
                final WritableLongChunk<Values> chunk = WritableLongChunk.makeWritableChunk(source.size());
                for (int ii = 0; ii < source.size(); ++ii) {
                    final long value = source.get(ii); // assume in nanoseconds
                    chunk.set(ii, (long)(value == QueryConstants.NULL_LONG ? QueryConstants.NULL_LONG : value / factor));
                }
                return chunk;
            }, ObjectChunk::getEmptyChunk, typeInfo.arrowField().isNullable());
        }
)

import org.apache.arrow.flight.CallHeaders
import org.apache.arrow.flight.FlightCallHeaders
import org.apache.arrow.flight.FlightClient
import org.apache.arrow.flight.HeaderCallOption
import org.apache.arrow.flight.Location
import org.apache.arrow.flight.client.ClientCookieMiddleware

allocator = new RootAllocator()
client = FlightClient.builder()
        .allocator(allocator)
        .location(Location.forGrpcInsecure("localhost", 10000))
        .intercept(new ClientCookieMiddleware.Factory())
        .build();
{
    CallHeaders headers = new FlightCallHeaders()
    headers.insert("Authorization", "Anonymous")
    // pre-shared key authentication would look like this:
    // headers.insert("Authorization", "io.deephaven.authentication.psk.PskAuthenticationHandler $PRE_SHARED_KEY_VALUE$")
    headers.insert("x-deephaven-auth-cookie-request", "true")
    client.handshake(new HeaderCallOption(headers))
}

import org.apache.arrow.vector.types.TimeUnit
import org.apache.arrow.vector.types.pojo.ArrowType
import org.apache.arrow.vector.types.pojo.FieldType
import org.apache.arrow.vector.types.pojo.Field
import org.apache.arrow.vector.types.pojo.Schema

// Prepare the schema with appropriate metadata to specify the target column type for the Deephaven table
COLUMN_NAME = "duration"
Map<String, String> attrs = new HashMap<>();
attrs.put("deephaven:type", "long");
// if this was an array/vector type, we would set the componentType, too:
// attrs.put("deephaven:componentType", "canonical.class.path.to.type");
boolean nullable = true
ArrowType arrowType = new ArrowType.Duration(TimeUnit.MILLISECOND)
FieldType fieldType = new FieldType(nullable, arrowType, null, attrs);
Schema schema = new Schema(Collections.singletonList(new Field(COLUMN_NAME, fieldType, null)));

// Upload data:
import org.apache.arrow.flight.FlightDescriptor
import org.apache.arrow.flight.AsyncPutListener
import org.apache.arrow.vector.DurationVector
import org.apache.arrow.vector.VectorSchemaRoot

NUM_ROWS = 32
try (final VectorSchemaRoot source = VectorSchemaRoot.create(schema, allocator)) {
    source.allocateNew()
    source.setRowCount(NUM_ROWS)
    DurationVector vector = (DurationVector) source.getVector(0)

    rnd = new java.util.Random()
    for (int ii = 0; ii < NUM_ROWS; ++ii) {
        // random millisecond durations
        if (ii % 8 == 7) {
            vector.setNull(ii)
        } else {
            vector.set(ii, rnd.nextInt(1000))
        }
    }

    // upload the table into a QueryScope variable named "uploaded_table"
    FlightDescriptor descriptor = FlightDescriptor.path("scope", "uploaded_table")
    AsyncPutListener listener = new AsyncPutListener()
    FlightClient.ClientStreamListener putStream = client.startPut(descriptor, source, listener)
    putStream.putNext()
    putStream.completed()
    // ensure that the server finishes publishing the result table to the QueryScope
    listener.getResult()
}

TableTools.show(uploaded_table.meta())

// print the first ten rows of the table
TableTools.show(uploaded_table)

import org.apache.arrow.flight.FlightStream
import org.apache.arrow.flight.Ticket

Ticket ticket = new Ticket("s/uploaded_table".getBytes(StandardCharsets.UTF_8));
try (final FlightStream stream = client.getStream(ticket)) {
    while (stream.next()) {
        println(stream.getRoot().getVector(0))
    }
}

[PT0.979S, PT0.384S, PT0.186S, PT0.531S, PT0.542S, PT0.44S, PT0.576S, null, PT0.683S, PT0.699S, ... PT0.868S, null, PT0.118S, PT0.003S, PT0.356S, PT0.907S, PT0.799S, PT0.248S, PT0.372S, null]