Handle restricted output columns in Arrow Page Source

Restricted columns need to be handled when using table valued functions in query.

Author: elbinpallimalilibm

presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowPageSource.java

@@ -22,9 +22,16 @@
 import org.apache.arrow.vector.FieldVector;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.OptionalInt;
+import java.util.stream.IntStream;
 
 import static com.facebook.plugin.arrow.ArrowErrorCode.ARROW_FLIGHT_CLIENT_ERROR;
+import static com.facebook.plugin.arrow.ArrowErrorCode.ARROW_INTERNAL_ERROR;
+import static com.facebook.presto.common.Utils.checkArgument;
 import static java.util.Objects.requireNonNull;
 
 public class ArrowPageSource
@@ -34,6 +41,8 @@ public class ArrowPageSource
     private final List<ArrowColumnHandle> columnHandles;
     private final ArrowBlockBuilder arrowBlockBuilder;
     private final ClientClosingFlightStream flightStreamAndClient;
+    private final Optional<List<ArrowColumnHandle>> columnHandlesInSplit;
+    private final Map<Integer, Integer> columnIndexToVectorIndexMap;
     private boolean completed;
     private int currentPosition;
 
@@ -49,6 +58,8 @@ public ArrowPageSource(
         requireNonNull(clientHandler, "clientHandler is null");
         this.arrowBlockBuilder = requireNonNull(arrowBlockBuilder, "arrowBlockBuilder is null");
         this.flightStreamAndClient = clientHandler.getFlightStream(connectorSession, split);
+        this.columnHandlesInSplit = requireNonNull(split.getColumns(), "split.getColumns is null");
+        columnIndexToVectorIndexMap = getColumnIndexToVectorIndexMap();
     }
 
     @Override
@@ -98,8 +109,11 @@ public Page getNextPage()
         // Create blocks from the loaded Arrow record batch
         List<Block> blocks = new ArrayList<>();
         List<FieldVector> vectors = flightStreamAndClient.getRoot().getFieldVectors();
+
+        columnHandlesInSplit.ifPresent(arrowColumnHandles -> checkArgument(vectors.size() == arrowColumnHandles.size(), "Number of field vectors is not the same as number of column handles in split"));
+
         for (int columnIndex = 0; columnIndex < columnHandles.size(); columnIndex++) {
-            FieldVector vector = vectors.get(columnIndex);
+            FieldVector vector = vectors.get(columnIndexToVectorIndexMap.get(columnIndex));
             Type type = columnHandles.get(columnIndex).getColumnType();
             Block block = arrowBlockBuilder.buildBlockFromFieldVector(vector, type, flightStreamAndClient.getDictionaryProvider());
             blocks.add(block);
@@ -122,4 +136,44 @@ public void close()
             throw new ArrowException(ARROW_FLIGHT_CLIENT_ERROR, e.getMessage(), e);
         }
     }
+
+    private int getVectorIndexForColumnHandleIndex(int columnHandleIndex)
+    {
+        int vectorIndex;
+        if (columnHandlesInSplit.isPresent()) {
+            // If the ArrowSplit defines the list of columns for the data in the split,
+            // get the vector to read by finding the index of the required column in this list.
+            OptionalInt index = IntStream
+                    .range(0, columnHandlesInSplit.get().size())
+                    .filter(k -> columnHandlesInSplit.get().get(k).getColumnName().equals(columnHandles.get(columnHandleIndex).getColumnName()))
+                    .findFirst();
+            if (index.isPresent()) {
+                vectorIndex = index.getAsInt();
+            }
+            else {
+                throw new ArrowException(ARROW_INTERNAL_ERROR, "Unable to find column " + columnHandles.get(columnHandleIndex).getColumnName() + " in the column handles given in split");
+            }
+        }
+        else {
+            // If the ArrowSplit does not define the list of columns for the data in the split,
+            // assume that the vector to read is in the same position as the column handle
+            vectorIndex = columnHandleIndex;
+        }
+
+        return vectorIndex;
+    }
+
+    private Map<Integer, Integer> getColumnIndexToVectorIndexMap()
+    {
+        Map<Integer, Integer> columnIndexToVectorIndexMap = new HashMap<>();
+        if (columnHandlesInSplit.isPresent()) {
+            for (int i = 0; i < columnHandles.size(); i++) {
+                columnIndexToVectorIndexMap.put(i, getVectorIndexForColumnHandleIndex(i));
+            }
+        }
+        else {
+            IntStream.range(0, columnHandles.size()).forEach(i -> columnIndexToVectorIndexMap.put(i, i));
+        }
+        return columnIndexToVectorIndexMap;
+    }
 }

presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowSplit.java

@@ -23,23 +23,27 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 
 public class ArrowSplit
         implements ConnectorSplit
 {
     private final String schemaName;
     private final String tableName;
     private final byte[] flightEndpointBytes;
+    private final Optional<List<ArrowColumnHandle>> columns;
 
     @JsonCreator
     public ArrowSplit(
             @JsonProperty("schemaName") @Nullable String schemaName,
             @JsonProperty("tableName") String tableName,
-            @JsonProperty("flightEndpointBytes") byte[] flightEndpointBytes)
+            @JsonProperty("flightEndpointBytes") byte[] flightEndpointBytes,
+            @JsonProperty("columns") Optional<List<ArrowColumnHandle>> columns)
     {
         this.schemaName = schemaName;
         this.tableName = tableName;
         this.flightEndpointBytes = flightEndpointBytes;
+        this.columns = columns;
     }
 
     @Override
@@ -77,4 +81,10 @@ public byte[] getFlightEndpointBytes()
     {
         return flightEndpointBytes;
     }
+
+    @JsonProperty
+    public Optional<List<ArrowColumnHandle>> getColumns()
+    {
+        return columns;
+    }
 }

presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowSplitManager.java

@@ -23,6 +23,7 @@
 import org.apache.arrow.flight.FlightInfo;
 
 import java.util.List;
+import java.util.Optional;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.Objects.requireNonNull;
@@ -49,8 +50,14 @@ public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHand
                 .map(info -> new ArrowSplit(
                         tableHandle.getSchema(),
                         tableHandle.getTable(),
-                        info.serialize().array()))
+                        info.serialize().array(),
+                        Optional.empty()))
                 .collect(toImmutableList());
         return new FixedSplitSource(splits);
     }
+
+    public BaseArrowFlightClientHandler getClientHandler()
+    {
+        return clientHandler;
+    }
 }

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java

@@ -62,13 +62,30 @@ public static DistributedQueryRunner createQueryRunner(int flightServerPort) thr
         return createQueryRunner(flightServerPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty(), Optional.empty());
     }
 
+    public static DistributedQueryRunner createQueryRunner(int flightServerPort, boolean testingWithTVF) throws Exception
+    {
+        return createQueryRunner(flightServerPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty(), Optional.empty(), testingWithTVF);
+    }
+
     public static DistributedQueryRunner createQueryRunner(
             int flightServerPort,
             Map<String, String> extraProperties,
             Map<String, String> coordinatorProperties,
             Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher,
             Optional<Boolean> mTLSEnabled)
             throws Exception
+    {
+        return createQueryRunner(flightServerPort, extraProperties, coordinatorProperties, externalWorkerLauncher, mTLSEnabled, false);
+    }
+
+    public static DistributedQueryRunner createQueryRunner(
+            int flightServerPort,
+            Map<String, String> extraProperties,
+            Map<String, String> coordinatorProperties,
+            Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher,
+            Optional<Boolean> mTLSEnabled,
+            boolean testingWithTVF)
+            throws Exception
     {
         Session session = testSessionBuilder()
                 .setCatalog(ARROW_FLIGHT_CATALOG)
@@ -87,7 +104,7 @@ public static DistributedQueryRunner createQueryRunner(
 
         try {
             boolean nativeExecution = externalWorkerLauncher.isPresent();
-            queryRunner.installPlugin(new TestingArrowFlightPlugin(nativeExecution));
+            queryRunner.installPlugin(new TestingArrowFlightPlugin(nativeExecution, testingWithTVF));
             Map<String, String> catalogProperties = ImmutableMap.of("arrow-flight.server.port", String.valueOf(flightServerPort));
 
             ImmutableMap.Builder<String, String> properties = ImmutableMap.<String, String>builder()

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java

@@ -53,7 +53,7 @@ public class TestArrowFlightQueries
         extends AbstractTestQueries
 {
     private static final Logger logger = Logger.get(TestArrowFlightQueries.class);
-    private int serverPort;
+    protected int serverPort;
     private RootAllocator allocator;
     private FlightServer server;
     private DistributedQueryRunner arrowFlightQueryRunner;

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueriesWithTVF.java

@@ -0,0 +1,44 @@
+/*
+ * 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 com.facebook.plugin.arrow;
+
+import com.facebook.presto.testing.MaterializedResult;
+import com.facebook.presto.testing.QueryRunner;
+import org.testng.annotations.Test;
+
+import static com.facebook.presto.common.type.IntegerType.INTEGER;
+import static com.facebook.presto.testing.MaterializedResult.resultBuilder;
+import static com.facebook.presto.testing.assertions.Assert.assertEquals;
+
+public class TestArrowFlightQueriesWithTVF
+        extends TestArrowFlightQueries
+{
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        serverPort = ArrowFlightQueryRunner.findUnusedPort();
+        return ArrowFlightQueryRunner.createQueryRunner(serverPort, true);
+    }
+
+    @Test
+    public void testQueryFunction()
+    {
+        MaterializedResult actualRow = computeActual("SELECT id from TABLE(system.query_function('SELECT name, id FROM tpch.member WHERE id = 1', 'name VARCHAR, id INTEGER'))");
+        MaterializedResult expectedRow = resultBuilder(getSession(), INTEGER)
+                .row(1)
+                .build();
+        assertEquals(actualRow, expectedRow);
+    }
+}

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowSplit.java

@@ -24,6 +24,7 @@
 import java.net.URISyntaxException;
 import java.nio.ByteBuffer;
 import java.util.List;
+import java.util.Optional;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
@@ -48,7 +49,7 @@ public void setUp()
         Location location = new Location("http://localhost:8080");
         flightEndpoint = new FlightEndpoint(ticket, location);
         // Instantiate ArrowSplit with mock data
-        arrowSplit = new ArrowSplit(schemaName, tableName, flightEndpoint.serialize().array());
+        arrowSplit = new ArrowSplit(schemaName, tableName, flightEndpoint.serialize().array(), Optional.empty());
     }
 
     @Test

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/testingConnector/PrimitiveToPrestoTypeMappings.java

@@ -0,0 +1,41 @@
+/*
+ * 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 com.facebook.plugin.arrow.testingConnector;
+
+import com.facebook.presto.common.type.IntegerType;
+import com.facebook.presto.common.type.Type;
+import com.facebook.presto.spi.PrestoException;
+
+import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType;
+import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
+
+public final class PrimitiveToPrestoTypeMappings
+{
+    private PrimitiveToPrestoTypeMappings()
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    public static Type fromPrimitiveToPrestoType(String dataType)
+    {
+        switch (dataType) {
+            case "INTEGER":
+                return IntegerType.INTEGER;
+
+            case "VARCHAR":
+                return createUnboundedVarcharType();
+        }
+        throw new PrestoException(NOT_SUPPORTED, "Unsupported datatype '" + dataType + "' in the selected table.");
+    }
+}

presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/testingConnector/QueryArrowTableHandle.java

@@ -0,0 +1,77 @@
+/*
+ * 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 com.facebook.plugin.arrow.testingConnector;
+
+import com.facebook.plugin.arrow.ArrowColumnHandle;
+import com.facebook.plugin.arrow.ArrowTableHandle;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import java.util.UUID;
+
+public class QueryArrowTableHandle
+        extends ArrowTableHandle
+{
+    private final String query;
+    private final List<ArrowColumnHandle> columns;
+
+    public QueryArrowTableHandle(
+            @JsonProperty String query,
+            @JsonProperty List<ArrowColumnHandle> columns)
+    {
+        super("schema-" + UUID.randomUUID(), "table-" + UUID.randomUUID());
+        this.query = query;
+        this.columns = Collections.unmodifiableList(columns);
+    }
+
+    @JsonProperty("query")
+    public String getQuery()
+    {
+        return query;
+    }
+
+    @JsonProperty("columns")
+    public List<ArrowColumnHandle> getColumns()
+    {
+        return columns;
+    }
+
+    @Override
+    public String toString()
+    {
+        return query + ":" + columns;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        QueryArrowTableHandle that = (QueryArrowTableHandle) o;
+        return Objects.equals(getSchema(), that.getSchema()) && Objects.equals(getTable(), that.getTable()) &&
+                Objects.equals(getQuery(), that.getQuery()) && Objects.equals(getColumns(), that.getColumns());
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(getSchema(), getTable(), query, columns);
+    }
+}