FlightStreamUtils.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.flight.sql.util;

import static java.util.Objects.isNull;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.channels.Channels;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.UInt1Vector;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.complex.DenseUnionVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.ipc.ReadChannel;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.util.Text;

public class FlightStreamUtils {

  public static List<List<String>> getResults(FlightStream stream) {
    final List<List<String>> results = new ArrayList<>();
    while (stream.next()) {
      try (final VectorSchemaRoot root = stream.getRoot()) {
        final long rowCount = root.getRowCount();
        for (int i = 0; i < rowCount; ++i) {
          results.add(new ArrayList<>());
        }

        root.getSchema()
            .getFields()
            .forEach(
                field -> {
                  try (final FieldVector fieldVector = root.getVector(field.getName())) {
                    if (fieldVector instanceof VarCharVector) {
                      final VarCharVector varcharVector = (VarCharVector) fieldVector;
                      for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
                        final Text data = varcharVector.getObject(rowIndex);
                        results.get(rowIndex).add(isNull(data) ? null : data.toString());
                      }
                    } else if (fieldVector instanceof IntVector) {
                      for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
                        Object data = fieldVector.getObject(rowIndex);
                        results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data));
                      }
                    } else if (fieldVector instanceof VarBinaryVector) {
                      final VarBinaryVector varbinaryVector = (VarBinaryVector) fieldVector;
                      for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
                        final byte[] data = varbinaryVector.getObject(rowIndex);
                        final String output;
                        try {
                          output =
                              isNull(data)
                                  ? null
                                  : MessageSerializer.deserializeSchema(
                                          new ReadChannel(
                                              Channels.newChannel(new ByteArrayInputStream(data))))
                                      .toJson();
                        } catch (final IOException e) {
                          throw new RuntimeException("Failed to deserialize schema", e);
                        }
                        results.get(rowIndex).add(output);
                      }
                    } else if (fieldVector instanceof DenseUnionVector) {
                      final DenseUnionVector denseUnionVector = (DenseUnionVector) fieldVector;
                      for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
                        final Object data = denseUnionVector.getObject(rowIndex);
                        results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data));
                      }
                    } else if (fieldVector instanceof ListVector) {
                      for (int i = 0; i < fieldVector.getValueCount(); i++) {
                        if (!fieldVector.isNull(i)) {
                          List<Text> elements =
                              (List<Text>) ((ListVector) fieldVector).getObject(i);
                          List<String> values = new ArrayList<>();

                          for (Text element : elements) {
                            values.add(element.toString());
                          }
                          results.get(i).add(values.toString());
                        }
                      }

                    } else if (fieldVector instanceof UInt4Vector) {
                      final UInt4Vector uInt4Vector = (UInt4Vector) fieldVector;
                      for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
                        final Object data = uInt4Vector.getObject(rowIndex);
                        results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data));
                      }
                    } else if (fieldVector instanceof UInt1Vector) {
                      final UInt1Vector uInt1Vector = (UInt1Vector) fieldVector;
                      for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
                        final Object data = uInt1Vector.getObject(rowIndex);
                        results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data));
                      }
                    } else if (fieldVector instanceof BitVector) {
                      for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
                        Object data = fieldVector.getObject(rowIndex);
                        results.get(rowIndex).add(isNull(data) ? null : Objects.toString(data));
                      }
                    } else {
                      throw new UnsupportedOperationException("Not yet implemented");
                    }
                  }
                });
      }
    }

    return results;
  }
}