options) {
+ return makeTable(new CaseInsensitiveStringMap(options));
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightArrowColumnVector.java b/src/main/java/org/apache/arrow/flight/spark/FlightArrowColumnVector.java
new file mode 100644
index 0000000..0d2af58
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightArrowColumnVector.java
@@ -0,0 +1,575 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.
+ */
+/*
+ * 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.spark;
+
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.DateDayVector;
+import org.apache.arrow.vector.DateMilliVector;
+import org.apache.arrow.vector.DecimalVector;
+import org.apache.arrow.vector.Float4Vector;
+import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.SmallIntVector;
+import org.apache.arrow.vector.TimeStampMicroVector;
+import org.apache.arrow.vector.TimeStampMicroTZVector;
+import org.apache.arrow.vector.TimeStampMilliVector;
+import org.apache.arrow.vector.TimeStampVector;
+import org.apache.arrow.vector.TinyIntVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.holders.NullableVarCharHolder;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.spark.sql.execution.arrow.FlightArrowUtils;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarArray;
+import org.apache.spark.sql.vectorized.ColumnarMap;
+import org.apache.spark.unsafe.types.UTF8String;
+
+
+/**
+ * A column vector backed by Apache Arrow. Currently calendar interval type and map type are not
+ * supported. This is a copy of ArrowColumnVector with added support for DateMilli and TimestampMilli
+ */
+public final class FlightArrowColumnVector extends ColumnVector {
+
+ private final ArrowVectorAccessor accessor;
+ private FlightArrowColumnVector[] childColumns;
+
+ @Override
+ public boolean hasNull() {
+ return accessor.getNullCount() > 0;
+ }
+
+ @Override
+ public int numNulls() {
+ return accessor.getNullCount();
+ }
+
+ @Override
+ public void close() {
+ if (childColumns != null) {
+ for (int i = 0; i < childColumns.length; i++) {
+ childColumns[i].close();
+ childColumns[i] = null;
+ }
+ childColumns = null;
+ }
+ accessor.close();
+ }
+
+ @Override
+ public boolean isNullAt(int rowId) {
+ return accessor.isNullAt(rowId);
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return accessor.getBoolean(rowId);
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return accessor.getByte(rowId);
+ }
+
+ @Override
+ public short getShort(int rowId) {
+ return accessor.getShort(rowId);
+ }
+
+ @Override
+ public int getInt(int rowId) {
+ return accessor.getInt(rowId);
+ }
+
+ @Override
+ public long getLong(int rowId) {
+ return accessor.getLong(rowId);
+ }
+
+ @Override
+ public float getFloat(int rowId) {
+ return accessor.getFloat(rowId);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return accessor.getDouble(rowId);
+ }
+
+ @Override
+ public Decimal getDecimal(int rowId, int precision, int scale) {
+ if (isNullAt(rowId)) {
+ return null;
+ }
+ return accessor.getDecimal(rowId, precision, scale);
+ }
+
+ @Override
+ public UTF8String getUTF8String(int rowId) {
+ if (isNullAt(rowId)) {
+ return null;
+ }
+ return accessor.getUTF8String(rowId);
+ }
+
+ @Override
+ public byte[] getBinary(int rowId) {
+ if (isNullAt(rowId)) {
+ return null;
+ }
+ return accessor.getBinary(rowId);
+ }
+
+ @Override
+ public ColumnarArray getArray(int rowId) {
+ if (isNullAt(rowId)) {
+ return null;
+ }
+ return accessor.getArray(rowId);
+ }
+
+ @Override
+ public ColumnarMap getMap(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public FlightArrowColumnVector getChild(int ordinal) {
+ return childColumns[ordinal];
+ }
+
+ public FlightArrowColumnVector(ValueVector vector) {
+ super(FlightArrowUtils.fromArrowField(vector.getField()));
+
+ if (vector instanceof BitVector) {
+ accessor = new BooleanAccessor((BitVector) vector);
+ } else if (vector instanceof TinyIntVector) {
+ accessor = new ByteAccessor((TinyIntVector) vector);
+ } else if (vector instanceof SmallIntVector) {
+ accessor = new ShortAccessor((SmallIntVector) vector);
+ } else if (vector instanceof IntVector) {
+ accessor = new IntAccessor((IntVector) vector);
+ } else if (vector instanceof BigIntVector) {
+ accessor = new LongAccessor((BigIntVector) vector);
+ } else if (vector instanceof Float4Vector) {
+ accessor = new FloatAccessor((Float4Vector) vector);
+ } else if (vector instanceof Float8Vector) {
+ accessor = new DoubleAccessor((Float8Vector) vector);
+ } else if (vector instanceof DecimalVector) {
+ accessor = new DecimalAccessor((DecimalVector) vector);
+ } else if (vector instanceof VarCharVector) {
+ accessor = new StringAccessor((VarCharVector) vector);
+ } else if (vector instanceof VarBinaryVector) {
+ accessor = new BinaryAccessor((VarBinaryVector) vector);
+ } else if (vector instanceof DateDayVector) {
+ accessor = new DateAccessor((DateDayVector) vector);
+ } else if (vector instanceof DateMilliVector) {
+ accessor = new DateMilliAccessor((DateMilliVector) vector);
+ } else if (vector instanceof TimeStampMicroVector) {
+ accessor = new TimestampMicroAccessor((TimeStampMicroVector) vector);
+ } else if (vector instanceof TimeStampMicroTZVector) {
+ accessor = new TimestampMicroTZAccessor((TimeStampMicroTZVector) vector);
+ } else if (vector instanceof TimeStampMilliVector) {
+ accessor = new TimestampMilliAccessor((TimeStampMilliVector) vector);
+ } else if (vector instanceof ListVector) {
+ ListVector listVector = (ListVector) vector;
+ accessor = new ArrayAccessor(listVector);
+ } else if (vector instanceof StructVector) {
+ StructVector structVector = (StructVector) vector;
+ accessor = new StructAccessor(structVector);
+
+ childColumns = new FlightArrowColumnVector[structVector.size()];
+ for (int i = 0; i < childColumns.length; ++i) {
+ childColumns[i] = new FlightArrowColumnVector(structVector.getVectorById(i));
+ }
+ } else {
+ System.out.println(vector);
+ throw new UnsupportedOperationException();
+ }
+ }
+
+ private abstract static class ArrowVectorAccessor {
+
+ private final ValueVector vector;
+
+ ArrowVectorAccessor(ValueVector vector) {
+ this.vector = vector;
+ }
+
+ // TODO: should be final after removing ArrayAccessor workaround
+ boolean isNullAt(int rowId) {
+ return vector.isNull(rowId);
+ }
+
+ final int getNullCount() {
+ return vector.getNullCount();
+ }
+
+ final void close() {
+ vector.close();
+ }
+
+ boolean getBoolean(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ byte getByte(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ short getShort(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ int getInt(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ long getLong(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ float getFloat(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ double getDouble(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ Decimal getDecimal(int rowId, int precision, int scale) {
+ throw new UnsupportedOperationException();
+ }
+
+ UTF8String getUTF8String(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ byte[] getBinary(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+
+ ColumnarArray getArray(int rowId) {
+ throw new UnsupportedOperationException();
+ }
+ }
+
+ private static class BooleanAccessor extends ArrowVectorAccessor {
+
+ private final BitVector accessor;
+
+ BooleanAccessor(BitVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final boolean getBoolean(int rowId) {
+ return accessor.get(rowId) == 1;
+ }
+ }
+
+ private static class ByteAccessor extends ArrowVectorAccessor {
+
+ private final TinyIntVector accessor;
+
+ ByteAccessor(TinyIntVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final byte getByte(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class ShortAccessor extends ArrowVectorAccessor {
+
+ private final SmallIntVector accessor;
+
+ ShortAccessor(SmallIntVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final short getShort(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class IntAccessor extends ArrowVectorAccessor {
+
+ private final IntVector accessor;
+
+ IntAccessor(IntVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final int getInt(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class LongAccessor extends ArrowVectorAccessor {
+
+ private final BigIntVector accessor;
+
+ LongAccessor(BigIntVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final long getLong(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class FloatAccessor extends ArrowVectorAccessor {
+
+ private final Float4Vector accessor;
+
+ FloatAccessor(Float4Vector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final float getFloat(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class DoubleAccessor extends ArrowVectorAccessor {
+
+ private final Float8Vector accessor;
+
+ DoubleAccessor(Float8Vector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final double getDouble(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class DecimalAccessor extends ArrowVectorAccessor {
+
+ private final DecimalVector accessor;
+
+ DecimalAccessor(DecimalVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final Decimal getDecimal(int rowId, int precision, int scale) {
+ if (isNullAt(rowId)) {
+ return null;
+ }
+ return Decimal.apply(accessor.getObject(rowId), precision, scale);
+ }
+ }
+
+ private static class StringAccessor extends ArrowVectorAccessor {
+
+ private final VarCharVector accessor;
+ private final NullableVarCharHolder stringResult = new NullableVarCharHolder();
+
+ StringAccessor(VarCharVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final UTF8String getUTF8String(int rowId) {
+ accessor.get(rowId, stringResult);
+ if (stringResult.isSet == 0) {
+ return null;
+ } else {
+ return UTF8String.fromAddress(null,
+ stringResult.buffer.memoryAddress() + stringResult.start,
+ stringResult.end - stringResult.start);
+ }
+ }
+ }
+
+ private static class BinaryAccessor extends ArrowVectorAccessor {
+
+ private final VarBinaryVector accessor;
+
+ BinaryAccessor(VarBinaryVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final byte[] getBinary(int rowId) {
+ return accessor.getObject(rowId);
+ }
+ }
+
+ private static class DateAccessor extends ArrowVectorAccessor {
+
+ private final DateDayVector accessor;
+
+ DateAccessor(DateDayVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final int getInt(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class DateMilliAccessor extends ArrowVectorAccessor {
+
+ private final DateMilliVector accessor;
+ private final double val = 1.0 / (24. * 60. * 60. * 1000.);
+
+ DateMilliAccessor(DateMilliVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final int getInt(int rowId) {
+ System.out.println(accessor.get(rowId) + " " + (accessor.get(rowId) * val) + " " + val);
+ return (int) (accessor.get(rowId) * val);
+ }
+ }
+
+ private static class TimestampMicroAccessor extends ArrowVectorAccessor {
+
+ private final TimeStampVector accessor;
+
+ TimestampMicroAccessor(TimeStampMicroVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final long getLong(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class TimestampMicroTZAccessor extends ArrowVectorAccessor {
+
+ private final TimeStampVector accessor;
+
+ TimestampMicroTZAccessor(TimeStampMicroTZVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final long getLong(int rowId) {
+ return accessor.get(rowId);
+ }
+ }
+
+ private static class TimestampMilliAccessor extends ArrowVectorAccessor {
+
+ private final TimeStampVector accessor;
+
+ TimestampMilliAccessor(TimeStampMilliVector vector) {
+ super(vector);
+ this.accessor = vector;
+ }
+
+ @Override
+ final long getLong(int rowId) {
+ return accessor.get(rowId) * 1000;
+ }
+ }
+
+ private static class ArrayAccessor extends ArrowVectorAccessor {
+
+ private final ListVector accessor;
+ private final FlightArrowColumnVector arrayData;
+
+ ArrayAccessor(ListVector vector) {
+ super(vector);
+ this.accessor = vector;
+ this.arrayData = new FlightArrowColumnVector(vector.getDataVector());
+ }
+
+ @Override
+ final boolean isNullAt(int rowId) {
+ // TODO: Workaround if vector has all non-null values, see ARROW-1948
+ if (accessor.getValueCount() > 0 && accessor.getValidityBuffer().capacity() == 0) {
+ return false;
+ } else {
+ return super.isNullAt(rowId);
+ }
+ }
+
+ @Override
+ final ColumnarArray getArray(int rowId) {
+ ArrowBuf offsets = accessor.getOffsetBuffer();
+ int index = rowId * ListVector.OFFSET_WIDTH;
+ int start = offsets.getInt(index);
+ int end = offsets.getInt(index + ListVector.OFFSET_WIDTH);
+ return new ColumnarArray(arrayData, start, end - start);
+ }
+ }
+
+ /**
+ * Any call to "get" method will throw UnsupportedOperationException.
+ *
+ * Access struct values in a ArrowColumnVector doesn't use this accessor. Instead, it uses
+ * getStruct() method defined in the parent class. Any call to "get" method in this class is a
+ * bug in the code.
+ */
+ private static class StructAccessor extends ArrowVectorAccessor {
+
+ StructAccessor(StructVector vector) {
+ super(vector);
+ }
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightClientFactory.java b/src/main/java/org/apache/arrow/flight/spark/FlightClientFactory.java
new file mode 100644
index 0000000..671b6ab
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightClientFactory.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.io.ByteArrayInputStream;
+import java.io.InputStream;
+
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+
+public class FlightClientFactory implements AutoCloseable {
+ private final BufferAllocator allocator = new RootAllocator();
+ private final Location defaultLocation;
+ private final FlightClientOptions clientOptions;
+
+ public FlightClientFactory(Location defaultLocation, FlightClientOptions clientOptions) {
+ this.defaultLocation = defaultLocation;
+ this.clientOptions = clientOptions;
+ }
+
+ public FlightClient apply() {
+ FlightClient.Builder builder = FlightClient.builder(allocator, defaultLocation);
+
+ if (!clientOptions.getTrustedCertificates().isEmpty()) {
+ builder.trustedCertificates(new ByteArrayInputStream(clientOptions.getTrustedCertificates().getBytes()));
+ }
+
+ if (!clientOptions.getClientCertificate().isEmpty()) {
+ InputStream clientCert = new ByteArrayInputStream(clientOptions.getClientCertificate().getBytes());
+ InputStream clientKey = new ByteArrayInputStream(clientOptions.getClientKey().getBytes());
+ builder.clientCertificate(clientCert, clientKey);
+ }
+
+ // Add client middleware
+ clientOptions.getMiddleware().stream().forEach(middleware -> builder.intercept(middleware));
+
+ FlightClient client = builder.build();
+ if (!clientOptions.getUsername().isEmpty()) {
+ client.authenticateBasic(clientOptions.getUsername(), clientOptions.getPassword());
+ }
+
+ return client;
+ }
+
+ @Override
+ public void close() {
+ allocator.close();
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightClientMiddlewareFactory.java b/src/main/java/org/apache/arrow/flight/spark/FlightClientMiddlewareFactory.java
new file mode 100644
index 0000000..538b289
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightClientMiddlewareFactory.java
@@ -0,0 +1,25 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.io.Serializable;
+
+import org.apache.arrow.flight.FlightClientMiddleware;
+
+public interface FlightClientMiddlewareFactory extends FlightClientMiddleware.Factory, Serializable {
+
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightClientOptions.java b/src/main/java/org/apache/arrow/flight/spark/FlightClientOptions.java
new file mode 100644
index 0000000..6bbae18
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightClientOptions.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.io.Serializable;
+import java.util.List;
+
+public class FlightClientOptions implements Serializable {
+ private final String username;
+ private final String password;
+ private final String trustedCertificates;
+ private final String clientCertificate;
+ private final String clientKey;
+ private final List middleware;
+
+ public FlightClientOptions(String username, String password, String trustedCertificates, String clientCertificate, String clientKey, List middleware) {
+ this.username = username;
+ this.password = password;
+ this.trustedCertificates = trustedCertificates;
+ this.clientCertificate = clientCertificate;
+ this.clientKey = clientKey;
+ this.middleware = middleware;
+ }
+
+ public String getUsername() {
+ return username;
+ }
+
+ public String getPassword() {
+ return password;
+ }
+
+ public String getTrustedCertificates() {
+ return trustedCertificates;
+ }
+
+ public String getClientCertificate() {
+ return clientCertificate;
+ }
+
+ public String getClientKey() {
+ return clientKey;
+ }
+
+ public List getMiddleware() {
+ return middleware;
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightColumnarPartitionReader.java b/src/main/java/org/apache/arrow/flight/spark/FlightColumnarPartitionReader.java
new file mode 100644
index 0000000..d130673
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightColumnarPartitionReader.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.io.IOException;
+
+import org.apache.spark.sql.connector.read.PartitionReader;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.util.AutoCloseables;
+import org.apache.spark.sql.vectorized.ColumnVector;
+
+public class FlightColumnarPartitionReader implements PartitionReader {
+ private final FlightClientFactory clientFactory;;
+ private final FlightClient client;
+ private final FlightStream stream;
+
+ public FlightColumnarPartitionReader(FlightClientOptions clientOptions, FlightPartition partition) {
+ // TODO - Should we handle multiple locations?
+ clientFactory = new FlightClientFactory(partition.getEndpoint().get().getLocations().get(0), clientOptions);
+ client = clientFactory.apply();
+ stream = client.getStream(partition.getEndpoint().get().getTicket());
+ }
+
+ // This is written this way because the Spark interface iterates in a different way.
+ // E.g., .next() -> .get() vs. .hasNext() -> .next()
+ @Override
+ public boolean next() throws IOException {
+ try {
+ return stream.next();
+ } catch (RuntimeException e) {
+ throw new IOException(e);
+ }
+ }
+
+ @Override
+ public ColumnarBatch get() {
+ ColumnarBatch batch = new ColumnarBatch(
+ stream.getRoot().getFieldVectors()
+ .stream()
+ .map(FlightArrowColumnVector::new)
+ .toArray(ColumnVector[]::new)
+ );
+ batch.setNumRows(stream.getRoot().getRowCount());
+ return batch;
+ }
+
+ @Override
+ public void close() throws IOException {
+ try {
+ AutoCloseables.close(stream, client, clientFactory);
+ } catch (Exception e) {
+ throw new IOException(e);
+ }
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightEndpointWrapper.java b/src/main/java/org/apache/arrow/flight/spark/FlightEndpointWrapper.java
new file mode 100644
index 0000000..10aa8e7
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightEndpointWrapper.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.net.URI;
+import java.util.ArrayList;
+import java.util.stream.Collectors;
+
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.Ticket;
+
+// This is needed for FlightEndpoint to be Serializable in spark.
+// org.apache.arrow.flight.FlightEndpoint is a POJO of Serializable types.
+// However if spark is using build-in serialization instead of Kyro then we must implement Serializable
+public class FlightEndpointWrapper implements Serializable {
+ private FlightEndpoint inner;
+
+ public FlightEndpointWrapper(FlightEndpoint inner) {
+ this.inner = inner;
+ }
+
+ public FlightEndpoint get() {
+ return inner;
+ }
+
+ private void writeObject(ObjectOutputStream out) throws IOException {
+ ArrayList locations = inner.getLocations().stream().map(location -> location.getUri()).collect(Collectors.toCollection(ArrayList::new));
+ out.writeObject(locations);
+ out.write(inner.getTicket().getBytes());
+ }
+
+ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+ @SuppressWarnings("unchecked")
+ Location[] locations = ((ArrayList) in.readObject()).stream().map(l -> new Location(l)).toArray(Location[]::new);
+ byte[] ticket = in.readAllBytes();
+ this.inner = new FlightEndpoint(new Ticket(ticket), locations);
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightPartition.java b/src/main/java/org/apache/arrow/flight/spark/FlightPartition.java
new file mode 100644
index 0000000..db837e8
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightPartition.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import org.apache.spark.sql.connector.read.InputPartition;
+
+public class FlightPartition implements InputPartition {
+ private final FlightEndpointWrapper endpoint;
+
+ public FlightPartition(FlightEndpointWrapper endpoint) {
+ this.endpoint = endpoint;
+ }
+
+ @Override
+ public String[] preferredLocations() {
+ return endpoint.get().getLocations().stream().map(location -> location.getUri().getHost()).toArray(String[]::new);
+ }
+
+ public FlightEndpointWrapper getEndpoint() {
+ return endpoint;
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightPartitionReader.java b/src/main/java/org/apache/arrow/flight/spark/FlightPartitionReader.java
new file mode 100644
index 0000000..631d26d
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightPartitionReader.java
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.Optional;
+
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.util.AutoCloseables;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.read.PartitionReader;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+public class FlightPartitionReader implements PartitionReader {
+ private final FlightClientFactory clientFactory;;
+ private final FlightClient client;
+ private final FlightStream stream;
+ private Optional> batch;
+ private InternalRow row;
+
+ public FlightPartitionReader(FlightClientOptions clientOptions, FlightPartition partition) {
+ // TODO - Should we handle multiple locations?
+ clientFactory = new FlightClientFactory(partition.getEndpoint().get().getLocations().get(0), clientOptions);
+ client = clientFactory.apply();
+ stream = client.getStream(partition.getEndpoint().get().getTicket());
+ }
+
+ private Iterator getNextBatch() {
+ ColumnarBatch batch = new ColumnarBatch(
+ stream.getRoot().getFieldVectors()
+ .stream()
+ .map(FlightArrowColumnVector::new)
+ .toArray(ColumnVector[]::new)
+ );
+ batch.setNumRows(stream.getRoot().getRowCount());
+ return batch.rowIterator();
+ }
+
+ // This is written this way because the Spark interface iterates in a different way.
+ // E.g., .next() -> .get() vs. .hasNext() -> .next()
+ @Override
+ public boolean next() throws IOException {
+ try {
+ // Try the iterator first then get next batch
+ // Not quite rust match expressions...
+ return batch.map(currentBatch -> {
+ // Are there still rows in this batch?
+ if (currentBatch.hasNext()) {
+ row = currentBatch.next();
+ return true;
+ // No more rows, get the next batch
+ } else {
+ // Is there another batch?
+ if (stream.next()) {
+ // Yes, then fetch it.
+ Iterator nextBatch = getNextBatch();
+ batch = Optional.of(nextBatch);
+ if (currentBatch.hasNext()) {
+ row = currentBatch.next();
+ return true;
+ // Odd, we got an empty batch
+ } else {
+ return false;
+ }
+ // This partition / stream is complete
+ } else {
+ return false;
+ }
+ }
+ // Fetch the first batch
+ }).orElseGet(() -> {
+ // Is the stream empty?
+ if (stream.next()) {
+ // No, then fetch the first batch
+ Iterator firstBatch = getNextBatch();
+ batch = Optional.of(firstBatch);
+ if (firstBatch.hasNext()) {
+ row = firstBatch.next();
+ return true;
+ // Odd, we got an empty batch
+ } else {
+ return false;
+ }
+ // The stream was empty...
+ } else {
+ return false;
+ }
+ });
+ } catch (RuntimeException e) {
+ throw new IOException(e);
+ }
+ }
+
+ @Override
+ public InternalRow get() {
+ return row;
+ }
+
+ @Override
+ public void close() throws IOException {
+ try {
+ AutoCloseables.close(stream, client, clientFactory);
+ } catch (Exception e) {
+ throw new IOException(e);
+ }
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightPartitionReaderFactory.java b/src/main/java/org/apache/arrow/flight/spark/FlightPartitionReaderFactory.java
new file mode 100644
index 0000000..addc952
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightPartitionReaderFactory.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.connector.read.InputPartition;
+import org.apache.spark.sql.connector.read.PartitionReader;
+import org.apache.spark.sql.connector.read.PartitionReaderFactory;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+public class FlightPartitionReaderFactory implements PartitionReaderFactory {
+ private final Broadcast clientOptions;
+
+ public FlightPartitionReaderFactory(Broadcast clientOptions) {
+ this.clientOptions = clientOptions;
+ }
+
+ @Override
+ public PartitionReader createReader(InputPartition iPartition) {
+ // This feels wrong but this is what upstream spark sources do to.
+ FlightPartition partition = (FlightPartition) iPartition;
+ return new FlightPartitionReader(clientOptions.getValue(), partition);
+ }
+
+ @Override
+ public PartitionReader createColumnarReader(InputPartition iPartition) {
+ // This feels wrong but this is what upstream spark sources do to.
+ FlightPartition partition = (FlightPartition) iPartition;
+ return new FlightColumnarPartitionReader(clientOptions.getValue(), partition);
+ }
+
+ @Override
+ public boolean supportColumnarReads(InputPartition partition) {
+ return true;
+ }
+
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightScan.java b/src/main/java/org/apache/arrow/flight/spark/FlightScan.java
new file mode 100644
index 0000000..48bfe9f
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightScan.java
@@ -0,0 +1,63 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import org.apache.spark.sql.connector.read.Scan;
+
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.sql.connector.read.Batch;
+import org.apache.spark.sql.connector.read.InputPartition;
+import org.apache.spark.sql.connector.read.PartitionReaderFactory;
+import org.apache.spark.sql.types.StructType;
+
+public class FlightScan implements Scan, Batch {
+ private final StructType schema;
+ private final FlightInfo info;
+ private final Broadcast clientOptions;
+
+ public FlightScan(StructType schema, FlightInfo info, Broadcast clientOptions) {
+ this.schema = schema;
+ this.info = info;
+ this.clientOptions = clientOptions;
+ }
+
+ @Override
+ public StructType readSchema() {
+ return schema;
+ }
+
+ @Override
+ public Batch toBatch() {
+ return this;
+ }
+
+ @Override
+ public InputPartition[] planInputPartitions() {
+ InputPartition[] batches = info.getEndpoints().stream().map(endpoint -> {
+ FlightEndpointWrapper endpointWrapper = new FlightEndpointWrapper(endpoint);
+ return new FlightPartition(endpointWrapper);
+ }).toArray(InputPartition[]::new);
+ return batches;
+ }
+
+ @Override
+ public PartitionReaderFactory createReaderFactory() {
+ return new FlightPartitionReaderFactory(clientOptions);
+ }
+
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightScanBuilder.java b/src/main/java/org/apache/arrow/flight/spark/FlightScanBuilder.java
new file mode 100644
index 0000000..6932db5
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightScanBuilder.java
@@ -0,0 +1,299 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.types.FloatingPointPrecision;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.sql.connector.read.Scan;
+import org.apache.spark.sql.connector.read.ScanBuilder;
+import org.apache.spark.sql.connector.read.SupportsPushDownFilters;
+import org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns;
+import org.apache.spark.sql.sources.*;
+import org.apache.spark.sql.types.*;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import scala.collection.JavaConversions;
+
+import com.google.common.collect.Lists;
+import com.google.common.base.Joiner;
+
+public class FlightScanBuilder implements ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownFilters {
+ private static final Logger LOGGER = LoggerFactory.getLogger(FlightScanBuilder.class);
+ private static final Joiner WHERE_JOINER = Joiner.on(" and ");
+ private static final Joiner PROJ_JOINER = Joiner.on(", ");
+ private SchemaResult flightSchema;
+ private StructType schema;
+ private final Location location;
+ private final Broadcast clientOptions;
+ private FlightDescriptor descriptor;
+ private String sql;
+ private Filter[] pushed;
+
+ public FlightScanBuilder(Location location, Broadcast clientOptions, String sql) {
+ this.location = location;
+ this.clientOptions = clientOptions;
+ this.sql = sql;
+ descriptor = getDescriptor(sql);
+ }
+
+ private class Client implements AutoCloseable {
+ private final FlightClientFactory clientFactory;
+ private final FlightClient client;
+
+ public Client(Location location, FlightClientOptions clientOptions) {
+ this.clientFactory = new FlightClientFactory(location, clientOptions);
+ this.client = clientFactory.apply();
+ }
+
+ public FlightClient get() {
+ return client;
+ }
+
+ @Override
+ public void close() throws Exception {
+ AutoCloseables.close(client, clientFactory);
+ }
+ }
+
+ private void getFlightSchema(FlightDescriptor descriptor) {
+ try (Client client = new Client(location, clientOptions.getValue())) {
+ LOGGER.info("getSchema() descriptor: %s", descriptor);
+ flightSchema = client.get().getSchema(descriptor);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ public Scan build() {
+ try (Client client = new Client(location, clientOptions.getValue())) {
+ FlightDescriptor descriptor = FlightDescriptor.command(sql.getBytes());
+ LOGGER.info("getInfo() descriptor: %s", descriptor);
+ FlightInfo info = client.get().getInfo(descriptor);
+ return new FlightScan(readSchema(), info, clientOptions);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private boolean canBePushed(Filter filter) {
+ if (filter instanceof IsNotNull) {
+ return true;
+ } else if (filter instanceof EqualTo) {
+ return true;
+ }
+ if (filter instanceof GreaterThan) {
+ return true;
+ }
+ if (filter instanceof GreaterThanOrEqual) {
+ return true;
+ }
+ if (filter instanceof LessThan) {
+ return true;
+ }
+ if (filter instanceof LessThanOrEqual) {
+ return true;
+ }
+ LOGGER.error("Cant push filter of type " + filter.toString());
+ return false;
+ }
+
+ private String valueToString(Object value) {
+ if (value instanceof String) {
+ return String.format("'%s'", value);
+ }
+ return value.toString();
+ }
+
+ private String generateWhereClause(List pushed) {
+ List filterStr = Lists.newArrayList();
+ for (Filter filter : pushed) {
+ if (filter instanceof IsNotNull) {
+ filterStr.add(String.format("isnotnull(\"%s\")", ((IsNotNull) filter).attribute()));
+ } else if (filter instanceof EqualTo) {
+ filterStr.add(String.format("\"%s\" = %s", ((EqualTo) filter).attribute(), valueToString(((EqualTo) filter).value())));
+ } else if (filter instanceof GreaterThan) {
+ filterStr.add(String.format("\"%s\" > %s", ((GreaterThan) filter).attribute(), valueToString(((GreaterThan) filter).value())));
+ } else if (filter instanceof GreaterThanOrEqual) {
+ filterStr.add(String.format("\"%s\" <= %s", ((GreaterThanOrEqual) filter).attribute(), valueToString(((GreaterThanOrEqual) filter).value())));
+ } else if (filter instanceof LessThan) {
+ filterStr.add(String.format("\"%s\" < %s", ((LessThan) filter).attribute(), valueToString(((LessThan) filter).value())));
+ } else if (filter instanceof LessThanOrEqual) {
+ filterStr.add(String.format("\"%s\" <= %s", ((LessThanOrEqual) filter).attribute(), valueToString(((LessThanOrEqual) filter).value())));
+ }
+ //todo fill out rest of Filter types
+ }
+ return WHERE_JOINER.join(filterStr);
+ }
+
+ private FlightDescriptor getDescriptor(String sql) {
+ return FlightDescriptor.command(sql.getBytes());
+ }
+
+ private void mergeWhereDescriptors(String whereClause) {
+ sql = String.format("select * from (%s) as where_merge where %s", sql, whereClause);
+ descriptor = getDescriptor(sql);
+ }
+
+ @Override
+ public Filter[] pushFilters(Filter[] filters) {
+ List notPushed = Lists.newArrayList();
+ List pushed = Lists.newArrayList();
+ for (Filter filter : filters) {
+ boolean isPushed = canBePushed(filter);
+ if (isPushed) {
+ pushed.add(filter);
+ } else {
+ notPushed.add(filter);
+ }
+ }
+ this.pushed = pushed.toArray(new Filter[0]);
+ if (!pushed.isEmpty()) {
+ String whereClause = generateWhereClause(pushed);
+ mergeWhereDescriptors(whereClause);
+ getFlightSchema(descriptor);
+ }
+ return notPushed.toArray(new Filter[0]);
+ }
+
+ @Override
+ public Filter[] pushedFilters() {
+ return pushed;
+ }
+
+ private DataType sparkFromArrow(FieldType fieldType) {
+ switch (fieldType.getType().getTypeID()) {
+ case Null:
+ return DataTypes.NullType;
+ case Struct:
+ throw new UnsupportedOperationException("have not implemented Struct type yet");
+ case List:
+ throw new UnsupportedOperationException("have not implemented List type yet");
+ case FixedSizeList:
+ throw new UnsupportedOperationException("have not implemented FixedSizeList type yet");
+ case Union:
+ throw new UnsupportedOperationException("have not implemented Union type yet");
+ case Int:
+ ArrowType.Int intType = (ArrowType.Int) fieldType.getType();
+ int bitWidth = intType.getBitWidth();
+ if (bitWidth == 8) {
+ return DataTypes.ByteType;
+ } else if (bitWidth == 16) {
+ return DataTypes.ShortType;
+ } else if (bitWidth == 32) {
+ return DataTypes.IntegerType;
+ } else if (bitWidth == 64) {
+ return DataTypes.LongType;
+ }
+ throw new UnsupportedOperationException("unknown int type with bitwidth " + bitWidth);
+ case FloatingPoint:
+ ArrowType.FloatingPoint floatType = (ArrowType.FloatingPoint) fieldType.getType();
+ FloatingPointPrecision precision = floatType.getPrecision();
+ switch (precision) {
+ case HALF:
+ case SINGLE:
+ return DataTypes.FloatType;
+ case DOUBLE:
+ return DataTypes.DoubleType;
+ }
+ case Utf8:
+ return DataTypes.StringType;
+ case Binary:
+ case FixedSizeBinary:
+ return DataTypes.BinaryType;
+ case Bool:
+ return DataTypes.BooleanType;
+ case Decimal:
+ throw new UnsupportedOperationException("have not implemented Decimal type yet");
+ case Date:
+ return DataTypes.DateType;
+ case Time:
+ return DataTypes.TimestampType; // note i don't know what this will do!
+ case Timestamp:
+ return DataTypes.TimestampType;
+ case Interval:
+ return DataTypes.CalendarIntervalType;
+ case NONE:
+ return DataTypes.NullType;
+ default:
+ throw new IllegalStateException("Unexpected value: " + fieldType);
+ }
+ }
+
+ private StructType readSchemaImpl() {
+ if (flightSchema == null) {
+ getFlightSchema(descriptor);
+ }
+ StructField[] fields = flightSchema.getSchema().getFields().stream()
+ .map(field -> new StructField(field.getName(),
+ sparkFromArrow(field.getFieldType()),
+ field.isNullable(),
+ Metadata.empty()))
+ .toArray(StructField[]::new);
+ return new StructType(fields);
+ }
+
+ public StructType readSchema() {
+ if (schema == null) {
+ schema = readSchemaImpl();
+ }
+ return schema;
+ }
+
+ private void mergeProjDescriptors(String projClause) {
+ sql = String.format("select %s from (%s) as proj_merge", projClause, sql);
+ descriptor = getDescriptor(sql);
+ }
+
+ @Override
+ public void pruneColumns(StructType requiredSchema) {
+ if (requiredSchema.toSeq().isEmpty()) {
+ return;
+ }
+ StructType schema = readSchema();
+ List fields = Lists.newArrayList();
+ List fieldsLeft = Lists.newArrayList();
+ Map fieldNames = JavaConversions.seqAsJavaList(schema.toSeq()).stream()
+ .collect(Collectors.toMap(StructField::name, f -> f));
+ for (StructField field : JavaConversions.seqAsJavaList(requiredSchema.toSeq())) {
+ String name = field.name();
+ StructField f = fieldNames.remove(name);
+ if (f != null) {
+ fields.add(String.format("\"%s\"", name));
+ fieldsLeft.add(f);
+ }
+ }
+ if (!fieldNames.isEmpty()) {
+ this.schema = new StructType(fieldsLeft.toArray(new StructField[0]));
+ mergeProjDescriptors(PROJ_JOINER.join(fields));
+ getFlightSchema(descriptor);
+ }
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightSparkContext.java b/src/main/java/org/apache/arrow/flight/spark/FlightSparkContext.java
new file mode 100644
index 0000000..f57a9af
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightSparkContext.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.sql.DataFrameReader;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSession;
+
+public class FlightSparkContext {
+
+ private SparkConf conf;
+
+ private final DataFrameReader reader;
+
+ public FlightSparkContext(SparkSession spark) {
+ this.conf = spark.sparkContext().getConf();
+ reader = spark.read().format("org.apache.arrow.flight.spark");
+ }
+
+ public Dataset read(String s) {
+ return reader.option("port", Integer.parseInt(conf.get("spark.flight.endpoint.port")))
+ .option("uri", String.format(
+ "grpc://%s:%s",
+ conf.get("spark.flight.endpoint.host"),
+ conf.get("spark.flight.endpoint.port")))
+ .option("username", conf.get("spark.flight.auth.username"))
+ .option("password", conf.get("spark.flight.auth.password"))
+ .load(s);
+ }
+
+ public Dataset readSql(String s) {
+ return reader.option("port", Integer.parseInt(conf.get("spark.flight.endpoint.port")))
+ .option("uri", String.format(
+ "grpc://%s:%s",
+ conf.get("spark.flight.endpoint.host"),
+ conf.get("spark.flight.endpoint.port")))
+ .option("username", conf.get("spark.flight.auth.username"))
+ .option("password", conf.get("spark.flight.auth.password"))
+ .load(s);
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/FlightTable.java b/src/main/java/org/apache/arrow/flight/spark/FlightTable.java
new file mode 100644
index 0000000..ee954ae
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/FlightTable.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.util.Set;
+
+import org.apache.arrow.flight.Location;
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.sql.connector.catalog.SupportsRead;
+import org.apache.spark.sql.connector.catalog.Table;
+import org.apache.spark.sql.connector.catalog.TableCapability;
+import org.apache.spark.sql.connector.read.ScanBuilder;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.util.CaseInsensitiveStringMap;
+
+public class FlightTable implements Table, SupportsRead {
+ private static final Set CAPABILITIES = Set.of(TableCapability.BATCH_READ);
+ private final String name;
+ private final Location location;
+ private final String sql;
+ private final Broadcast clientOptions;
+ private StructType schema;
+
+ public FlightTable(String name, Location location, String sql, Broadcast clientOptions) {
+ this.name = name;
+ this.location = location;
+ this.sql = sql;
+ this.clientOptions = clientOptions;
+ }
+
+ @Override
+ public String name() {
+ return name;
+ }
+
+ @Override
+ public StructType schema() {
+ if (schema == null) {
+ FlightScanBuilder scanBuilder = new FlightScanBuilder(location, clientOptions, sql);
+ schema = scanBuilder.readSchema();
+ }
+ return schema;
+ }
+
+ // TODO - We could probably implement partitioning() but it would require server side support
+
+ @Override
+ public Set capabilities() {
+ // We only support reading for now
+ return CAPABILITIES;
+ }
+
+ @Override
+ public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) {
+ return new FlightScanBuilder(location, clientOptions, sql);
+ }
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/TokenClientMiddleware.java b/src/main/java/org/apache/arrow/flight/spark/TokenClientMiddleware.java
new file mode 100644
index 0000000..8055c64
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/TokenClientMiddleware.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import org.apache.arrow.flight.CallHeaders;
+import org.apache.arrow.flight.CallStatus;
+import org.apache.arrow.flight.FlightClientMiddleware;
+
+public class TokenClientMiddleware implements FlightClientMiddleware {
+ private final String token;
+
+ public TokenClientMiddleware(String token) {
+ this.token = token;
+ }
+
+ @Override
+ public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
+ outgoingHeaders.insert("authorization", String.format("Bearer %s", token));
+ }
+
+ @Override
+ public void onHeadersReceived(CallHeaders incomingHeaders) {
+ // Nothing needed here
+ }
+
+ @Override
+ public void onCallCompleted(CallStatus status) {
+ // Nothing needed here
+ }
+
+}
diff --git a/src/main/java/org/apache/arrow/flight/spark/TokenClientMiddlewareFactory.java b/src/main/java/org/apache/arrow/flight/spark/TokenClientMiddlewareFactory.java
new file mode 100644
index 0000000..c424030
--- /dev/null
+++ b/src/main/java/org/apache/arrow/flight/spark/TokenClientMiddlewareFactory.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import org.apache.arrow.flight.CallInfo;
+import org.apache.arrow.flight.FlightClientMiddleware;
+
+public class TokenClientMiddlewareFactory implements FlightClientMiddlewareFactory {
+ private final String token;
+
+ public TokenClientMiddlewareFactory(String token) {
+ this.token = token;
+ }
+
+ @Override
+ public FlightClientMiddleware onCallStarted(CallInfo info) {
+ return new TokenClientMiddleware(token);
+ }
+
+}
diff --git a/src/main/scala/org/apache/spark/sql/execution/arrow/FlightArrowUtils.scala b/src/main/scala/org/apache/spark/sql/execution/arrow/FlightArrowUtils.scala
new file mode 100644
index 0000000..f6c569c
--- /dev/null
+++ b/src/main/scala/org/apache/spark/sql/execution/arrow/FlightArrowUtils.scala
@@ -0,0 +1,135 @@
+/**
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark.sql.execution.arrow
+
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
+import org.apache.arrow.vector.types.{DateUnit, FloatingPointPrecision, TimeUnit}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import scala.collection.JavaConverters._
+
+/**
+ * FlightArrowUtils is a copy of ArrowUtils with extra support for DateMilli and TimestampMilli
+ */
+object FlightArrowUtils {
+
+ val rootAllocator = new RootAllocator(Long.MaxValue)
+
+ // todo: support more types.
+
+ /** Maps data type from Spark to Arrow. NOTE: timeZoneId required for TimestampTypes */
+ def toArrowType(dt: DataType, timeZoneId: String): ArrowType = dt match {
+ case BooleanType => ArrowType.Bool.INSTANCE
+ case ByteType => new ArrowType.Int(8, true)
+ case ShortType => new ArrowType.Int(8 * 2, true)
+ case IntegerType => new ArrowType.Int(8 * 4, true)
+ case LongType => new ArrowType.Int(8 * 8, true)
+ case FloatType => new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)
+ case DoubleType => new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)
+ case StringType => ArrowType.Utf8.INSTANCE
+ case BinaryType => ArrowType.Binary.INSTANCE
+ case DecimalType.Fixed(precision, scale) => new ArrowType.Decimal(precision, scale)
+ case DateType => new ArrowType.Date(DateUnit.DAY)
+ case TimestampType =>
+ if (timeZoneId == null) {
+ throw new UnsupportedOperationException(
+ s"${TimestampType.catalogString} must supply timeZoneId parameter")
+ } else {
+ new ArrowType.Timestamp(TimeUnit.MICROSECOND, timeZoneId)
+ }
+ case _ =>
+ throw new UnsupportedOperationException(s"Unsupported data type: ${dt.catalogString}")
+ }
+
+ def fromArrowType(dt: ArrowType): DataType = dt match {
+ case ArrowType.Bool.INSTANCE => BooleanType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 => ByteType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 2 => ShortType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 4 => IntegerType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 8 => LongType
+ case float: ArrowType.FloatingPoint
+ if float.getPrecision() == FloatingPointPrecision.SINGLE => FloatType
+ case float: ArrowType.FloatingPoint
+ if float.getPrecision() == FloatingPointPrecision.DOUBLE => DoubleType
+ case ArrowType.Utf8.INSTANCE => StringType
+ case ArrowType.Binary.INSTANCE => BinaryType
+ case d: ArrowType.Decimal => DecimalType(d.getPrecision, d.getScale)
+ case date: ArrowType.Date if date.getUnit == DateUnit.DAY || date.getUnit == DateUnit.MILLISECOND => DateType
+ case ts: ArrowType.Timestamp if ts.getUnit == TimeUnit.MICROSECOND || ts.getUnit == TimeUnit.MILLISECOND => TimestampType
+ case _ => throw new UnsupportedOperationException(s"Unsupported data type: $dt")
+ }
+
+ /** Maps field from Spark to Arrow. NOTE: timeZoneId required for TimestampType */
+ def toArrowField(
+ name: String, dt: DataType, nullable: Boolean, timeZoneId: String): Field = {
+ dt match {
+ case ArrayType(elementType, containsNull) =>
+ val fieldType = new FieldType(nullable, ArrowType.List.INSTANCE, null)
+ new Field(name, fieldType,
+ Seq(toArrowField("element", elementType, containsNull, timeZoneId)).asJava)
+ case StructType(fields) =>
+ val fieldType = new FieldType(nullable, ArrowType.Struct.INSTANCE, null)
+ new Field(name, fieldType,
+ fields.map { field =>
+ toArrowField(field.name, field.dataType, field.nullable, timeZoneId)
+ }.toSeq.asJava)
+ case dataType =>
+ val fieldType = new FieldType(nullable, toArrowType(dataType, timeZoneId), null)
+ new Field(name, fieldType, Seq.empty[Field].asJava)
+ }
+ }
+
+ def fromArrowField(field: Field): DataType = {
+ field.getType match {
+ case ArrowType.List.INSTANCE =>
+ val elementField = field.getChildren().get(0)
+ val elementType = fromArrowField(elementField)
+ ArrayType(elementType, containsNull = elementField.isNullable)
+ case ArrowType.Struct.INSTANCE =>
+ val fields = field.getChildren().asScala.map { child =>
+ val dt = fromArrowField(child)
+ StructField(child.getName, dt, child.isNullable)
+ }
+ StructType(fields)
+ case arrowType => fromArrowType(arrowType)
+ }
+ }
+
+ /** Maps schema from Spark to Arrow. NOTE: timeZoneId required for TimestampType in StructType */
+ def toArrowSchema(schema: StructType, timeZoneId: String): Schema = {
+ new Schema(schema.map { field =>
+ toArrowField(field.name, field.dataType, field.nullable, timeZoneId)
+ }.asJava)
+ }
+
+ def fromArrowSchema(schema: Schema): StructType = {
+ StructType(schema.getFields.asScala.map { field =>
+ val dt = fromArrowField(field)
+ StructField(field.getName, dt, field.isNullable)
+ })
+ }
+
+ /** Return Map with conf settings to be used in ArrowPythonRunner */
+ def getPythonRunnerConfMap(conf: SQLConf): Map[String, String] = {
+ val timeZoneConf = Seq(SQLConf.SESSION_LOCAL_TIMEZONE.key ->
+ conf.sessionLocalTimeZone)
+ val pandasColsByName = Seq(SQLConf.PANDAS_GROUPED_MAP_ASSIGN_COLUMNS_BY_NAME.key ->
+ conf.pandasGroupedMapAssignColumnsByName.toString)
+ Map(timeZoneConf ++ pandasColsByName: _*)
+ }
+}
diff --git a/src/test/java/com/dremio/spark/TestConnector.java b/src/test/java/com/dremio/spark/TestConnector.java
deleted file mode 100644
index 5889986..0000000
--- a/src/test/java/com/dremio/spark/TestConnector.java
+++ /dev/null
@@ -1,48 +0,0 @@
-package com.dremio.spark;
-
-import com.dremio.BaseTestQuery;
-import com.dremio.exec.ExecTest;
-import com.dremio.service.InitializerRegistry;
-import org.apache.spark.SparkConf;
-import org.apache.spark.sql.SQLContext;
-import org.junit.AfterClass;
-import org.junit.BeforeClass;
-import org.junit.Ignore;
-import org.junit.Test;
-
-import org.apache.spark.api.java.JavaSparkContext;
-
-public class TestConnector extends BaseTestQuery {
- private static SparkConf conf;
- private static JavaSparkContext sc;
- private static FlightSparkContext csc;
- private static InitializerRegistry registry;
-
- @BeforeClass
- public static void setUp() throws Exception {
- registry = new InitializerRegistry(ExecTest.CLASSPATH_SCAN_RESULT, getBindingProvider());
- registry.start();
- conf = new SparkConf()
- .setAppName("flightTest")
- .setMaster("local[*]")
-// .set("spark.driver.allowMultipleContexts","true")
- .set("spark.flight.endpoint.host", "localhost")
- .set("spark.flight.endpoint.port", "47470")
- .set("spark.flight.username", "dremio")
- .set("spark.flight.password", "dremio123")
- ;
- sc = new JavaSparkContext(conf);
- csc = FlightSparkContext.flightContext(sc);
- }
-
- @AfterClass
- public static void tearDown() throws Exception {
- registry.close();
- sc.close();
- }
-
- @Test
- public void testConnect() {
- csc.read("sys.options");
- }
-}
diff --git a/src/test/java/org/apache/arrow/flight/spark/TestConnector.java b/src/test/java/org/apache/arrow/flight/spark/TestConnector.java
new file mode 100644
index 0000000..4ac7c37
--- /dev/null
+++ b/src/test/java/org/apache/arrow/flight/spark/TestConnector.java
@@ -0,0 +1,307 @@
+/*
+ * Copyright (C) 2019 The flight-spark-source Authors
+ *
+ * Licensed 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.spark;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectOutputStream;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Consumer;
+
+import org.apache.arrow.flight.Action;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightServer;
+import org.apache.arrow.flight.FlightTestUtil;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.NoOpFlightProducer;
+import org.apache.arrow.flight.Result;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.auth.ServerAuthHandler;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.Text;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSession;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.Test.None;
+
+import com.google.common.collect.ImmutableList;
+
+public class TestConnector {
+ private static final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ private static Location location;
+ private static FlightServer server;
+ private static SparkSession spark;
+ private static FlightSparkContext csc;
+
+ @BeforeClass
+ public static void setUp() throws Exception {
+ server = FlightTestUtil.getStartedServer(location -> FlightServer.builder(allocator, location, new TestProducer()).authHandler(
+ new ServerAuthHandler() {
+ @Override
+ public Optional isValid(byte[] token) {
+ return Optional.of("xxx");
+ }
+
+ @Override
+ public boolean authenticate(ServerAuthSender outgoing, Iterator incoming) {
+ incoming.next();
+ outgoing.send(new byte[0]);
+ return true;
+ }
+ }).build()
+ );
+ location = server.getLocation();
+ spark = SparkSession.builder()
+ .appName("flightTest")
+ .master("local[*]")
+ .config("spark.driver.host", "127.0.0.1")
+ .config("spark.driver.allowMultipleContexts", "true")
+ .config("spark.flight.endpoint.host", location.getUri().getHost())
+ .config("spark.flight.endpoint.port", Integer.toString(location.getUri().getPort()))
+ .config("spark.flight.auth.username", "xxx")
+ .config("spark.flight.auth.password", "yyy")
+ .getOrCreate();
+ csc = new FlightSparkContext(spark);
+ }
+
+ @AfterClass
+ public static void tearDown() throws Exception {
+ AutoCloseables.close(server, allocator, spark);
+ }
+
+ private class DummyObjectOutputStream extends ObjectOutputStream {
+ public DummyObjectOutputStream() throws IOException {
+ super(new ByteArrayOutputStream());
+ }
+ }
+
+ @Test(expected = None.class)
+ public void testFlightPartitionReaderFactorySerialization() throws IOException {
+ List middleware = new ArrayList<>();
+ FlightClientOptions clientOptions = new FlightClientOptions("xxx", "yyy", "FooBar", "FooBar", "FooBar", middleware);
+ FlightPartitionReaderFactory readerFactory = new FlightPartitionReaderFactory(JavaSparkContext.fromSparkContext(spark.sparkContext()).broadcast(clientOptions));
+
+ try (ObjectOutputStream oos = new DummyObjectOutputStream()) {
+ oos.writeObject(readerFactory);
+ }
+ }
+
+ @Test(expected = None.class)
+ public void testFlightPartitionSerialization() throws IOException {
+ Ticket ticket = new Ticket("FooBar".getBytes());
+ FlightEndpoint endpoint = new FlightEndpoint(ticket, location);
+ FlightPartition partition = new FlightPartition(new FlightEndpointWrapper(endpoint));
+ try (ObjectOutputStream oos = new DummyObjectOutputStream()) {
+ oos.writeObject(partition);
+ }
+ }
+
+ @Test
+ public void testConnect() {
+ csc.read("test.table");
+ }
+
+ @Test
+ public void testRead() {
+ long count = csc.read("test.table").count();
+ Assert.assertEquals(20, count);
+ }
+
+ @Test
+ public void testSql() {
+ long count = csc.readSql("select * from test.table").count();
+ Assert.assertEquals(20, count);
+ }
+
+ @Test
+ public void testFilter() {
+ Dataset df = csc.readSql("select * from test.table");
+ long count = df.filter(df.col("symbol").equalTo("USDCAD")).count();
+ long countOriginal = csc.readSql("select * from test.table").count();
+ Assert.assertTrue(count < countOriginal);
+ }
+
+ private static class SizeConsumer implements Consumer {
+ private int length = 0;
+ private int width = 0;
+
+ @Override
+ public void accept(Row row) {
+ length += 1;
+ width = row.length();
+ }
+ }
+
+ @Test
+ public void testProject() {
+ Dataset df = csc.readSql("select * from test.table");
+ SizeConsumer c = new SizeConsumer();
+ df.select("bid", "ask", "symbol").toLocalIterator().forEachRemaining(c);
+ long count = c.width;
+ long countOriginal = csc.readSql("select * from test.table").columns().length;
+ Assert.assertTrue(count < countOriginal);
+ }
+
+ private static class TestProducer extends NoOpFlightProducer {
+ private boolean parallel = false;
+
+ @Override
+ public void doAction(CallContext context, Action action, StreamListener listener) {
+ parallel = true;
+ listener.onNext(new Result("ok".getBytes()));
+ listener.onCompleted();
+ }
+
+ @Override
+ public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) {
+ Schema schema;
+ List endpoints;
+ if (parallel) {
+ endpoints = ImmutableList.of(new FlightEndpoint(new Ticket(descriptor.getCommand()), location),
+ new FlightEndpoint(new Ticket(descriptor.getCommand()), location));
+ } else {
+ endpoints = ImmutableList.of(new FlightEndpoint(new Ticket(descriptor.getCommand()), location));
+ }
+ if (new String(descriptor.getCommand()).equals("select \"bid\", \"ask\", \"symbol\" from (select * from test.table))")) {
+ schema = new Schema(ImmutableList.of(
+ Field.nullable("bid", Types.MinorType.FLOAT8.getType()),
+ Field.nullable("ask", Types.MinorType.FLOAT8.getType()),
+ Field.nullable("symbol", Types.MinorType.VARCHAR.getType()))
+ );
+
+ } else {
+ schema = new Schema(ImmutableList.of(
+ Field.nullable("bid", Types.MinorType.FLOAT8.getType()),
+ Field.nullable("ask", Types.MinorType.FLOAT8.getType()),
+ Field.nullable("symbol", Types.MinorType.VARCHAR.getType()),
+ Field.nullable("bidsize", Types.MinorType.BIGINT.getType()),
+ Field.nullable("asksize", Types.MinorType.BIGINT.getType()))
+ );
+ }
+ return new FlightInfo(schema, descriptor, endpoints, 1000000, 10);
+ }
+
+ @Override
+ public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
+ final int size = (new String(ticket.getBytes()).contains("USDCAD")) ? 5 : 10;
+
+ if (new String(ticket.getBytes()).equals("select \"bid\", \"ask\", \"symbol\" from (select * from test.table))")) {
+ Float8Vector b = new Float8Vector("bid", allocator);
+ Float8Vector a = new Float8Vector("ask", allocator);
+ VarCharVector s = new VarCharVector("symbol", allocator);
+
+ VectorSchemaRoot root = VectorSchemaRoot.of(b, a, s);
+ listener.start(root);
+
+ //batch 1
+ root.allocateNew();
+ for (int i = 0; i < size; i++) {
+ b.set(i, (double) i);
+ a.set(i, (double) i);
+ s.set(i, (i % 2 == 0) ? new Text("USDCAD") : new Text("EURUSD"));
+ }
+ b.setValueCount(size);
+ a.setValueCount(size);
+ s.setValueCount(size);
+ root.setRowCount(size);
+ listener.putNext();
+
+ // batch 2
+
+ root.allocateNew();
+ for (int i = 0; i < size; i++) {
+ b.set(i, (double) i);
+ a.set(i, (double) i);
+ s.set(i, (i % 2 == 0) ? new Text("USDCAD") : new Text("EURUSD"));
+ }
+ b.setValueCount(size);
+ a.setValueCount(size);
+ s.setValueCount(size);
+ root.setRowCount(size);
+ listener.putNext();
+ root.clear();
+ listener.completed();
+ } else {
+ BigIntVector bs = new BigIntVector("bidsize", allocator);
+ BigIntVector as = new BigIntVector("asksize", allocator);
+ Float8Vector b = new Float8Vector("bid", allocator);
+ Float8Vector a = new Float8Vector("ask", allocator);
+ VarCharVector s = new VarCharVector("symbol", allocator);
+
+ VectorSchemaRoot root = VectorSchemaRoot.of(b, a, s, bs, as);
+ listener.start(root);
+
+ //batch 1
+ root.allocateNew();
+ for (int i = 0; i < size; i++) {
+ bs.set(i, (long) i);
+ as.set(i, (long) i);
+ b.set(i, (double) i);
+ a.set(i, (double) i);
+ s.set(i, (i % 2 == 0) ? new Text("USDCAD") : new Text("EURUSD"));
+ }
+ bs.setValueCount(size);
+ as.setValueCount(size);
+ b.setValueCount(size);
+ a.setValueCount(size);
+ s.setValueCount(size);
+ root.setRowCount(size);
+ listener.putNext();
+
+ // batch 2
+
+ root.allocateNew();
+ for (int i = 0; i < size; i++) {
+ bs.set(i, (long) i);
+ as.set(i, (long) i);
+ b.set(i, (double) i);
+ a.set(i, (double) i);
+ s.set(i, (i % 2 == 0) ? new Text("USDCAD") : new Text("EURUSD"));
+ }
+ bs.setValueCount(size);
+ as.setValueCount(size);
+ b.setValueCount(size);
+ a.setValueCount(size);
+ s.setValueCount(size);
+ root.setRowCount(size);
+ listener.putNext();
+ root.clear();
+ listener.completed();
+ }
+ }
+
+
+ }
+}
diff --git a/src/test/resources/logback-test.xml b/src/test/resources/logback-test.xml
index 4a54f7d..62e5964 100644
--- a/src/test/resources/logback-test.xml
+++ b/src/test/resources/logback-test.xml
@@ -1,13 +1,13 @@