Arrow Flight Connector Template#24427

Merged

BryanCutler merged 1 commit intoprestodb:masterfrom

BryanCutler:arrow-connector-java-23032

Jan 29, 2025

Contributor

BryanCutler commented Jan 24, 2025 •

edited

Loading

Description

This adds the Arrow Flight base module as a template to build connectors that use a Flight service to transfer data with Presto in Arrow format. A concrete Arrow Flight connector implementation will extend the BaseArrowFlightClientHandler that can connect to a Flight service that handles the required RPC calls from the client handler. An example connector implementation is provided in the unit testing.

Supersedes: #23032

RFC: https://github.com/prestodb/rfcs/blob/main/RFC-0004-arrow-flight-connector.md

Motivation and Context

Adds the ability to easily create an Arrow Flight connector that can connect with any data source supported by the Flight service and transfer data efficiently using the Arrow format.

Impact

Adding new base-arrow-flight module.

Test Plan

Unit tests added with testing Arrow Flight connector implementation. These include extension of com.facebook.presto.tests.AbstractTestQueries that test general queries over Flight. Additional roundtrip testing of all data types supported by the ArrowBlockBuilder is done. An ArrowFlightQueryRunner is also provided that can be run standalone for manual query testing.

Contributor checklist

Please make sure your submission complies with our contributing guide, in particular code style and commit standards.
PR description addresses the issue accurately and concisely. If the change is non-trivial, a GitHub Issue is referenced.
Documented new properties (with its default value), SQL syntax, functions, or other functionality.
If release notes are required, they follow the release notes guidelines.
Adequate tests were added if applicable.
CI passed.

Release Notes

Please follow release notes guidelines and fill in the release notes below.

== RELEASE NOTES ==

General Changes
* Add Arrow Flight connector :pr:`24427`
* Add documentation for the :doc:`/connector/base-arrow-flight`  :pr:`24427`

BryanCutler requested review from a team, elharo and steveburnett as code owners

January 24, 2025 19:18

BryanCutler requested a review from presto-oss

January 24, 2025 19:18

prestodb-ci added the from:IBM label

prestodb-ci requested review from a team, aditi-pandit and psnv03 and removed request for a team

January 24, 2025 19:18

BryanCutler mentioned this pull request

Arrow-flight connector #23032

Closed

6 tasks

BryanCutler requested a review from tdcmeehan

January 24, 2025 19:20

BryanCutler commented

View reviewed changes

presto-base-arrow-flight/pom.xml

+                          <groupId>org.apache.arrow</groupId>
+                          <artifactId>arrow-jdbc</artifactId>
+                          <version>${dep.arrow.version}</version>
+                          <scope>test</scope>

Contributor Author

BryanCutler Jan 24, 2025

moved this to test dependency

BryanCutler commented

View reviewed changes

presto-base-arrow-flight/pom.xml

+                      <dependency>
+                          <groupId>org.apache.arrow</groupId>
+                          <artifactId>arrow-memory-netty</artifactId>
+                          <scope>runtime</scope>

Contributor Author

BryanCutler Jan 24, 2025

This is needed as a runtime dependency

BryanCutler commented

View reviewed changes

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

+                  public Block buildBlockFromFieldVector(FieldVector vector, Type type, DictionaryProvider dictionaryProvider)
+                  {
+                      // Use Arrow dictionary to create a DictionaryBlock
+                      if (dictionaryProvider != null && vector.getField().getDictionary() != null) {

Contributor Author

BryanCutler Jan 24, 2025

A DictionaryBlock is created only if a dictionary provider is supplied (by the FlightStream) and the Arrow Field contains a dictionary

BryanCutler commented

View reviewed changes

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

+                              // TODO: need method handles to construct MapType
+                              throw new UnsupportedOperationException("MapType is currently unsupported");
+                          }
+                          case Struct: {

Contributor Author

BryanCutler Jan 24, 2025 •

edited

Loading

Added Struct -> RowType

BryanCutler commented

View reviewed changes

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

+                              return BooleanType.BOOLEAN;
+                          case Time:
+                              return TimeType.TIME;
+                          case List: {

Contributor Author

BryanCutler Jan 24, 2025

Added List type -> ArrayType

BryanCutler commented

View reviewed changes

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

+                              Field entryField = children.get(0);
+                              checkArgument(entryField.getChildren().size() == 2, "Arrow Map entries expected to have 2 child Fields, got: " + children.size());
+                              // TODO: need method handles to construct MapType
+                              throw new UnsupportedOperationException("MapType is currently unsupported");

Contributor Author

BryanCutler Jan 24, 2025

I could not implement the MapType because need to figure out how to define the MethodHandles. Any ideas for this @tdcmeehan ?

BryanCutler commented

View reviewed changes

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

+                      }
+                  }
+                  private void assignBlockFromValueVector(ValueVector vector, Type type, BlockBuilder builder, int startIndex, int endIndex)

Contributor Author

BryanCutler Jan 24, 2025

Modified this to work with nested data types

BryanCutler commented

View reviewed changes

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

+                      }
+                  }
+                  public void assignBlockFromListVector(ListVector vector, Type type, BlockBuilder builder, int startIndex, int endIndex)

Contributor Author

BryanCutler Jan 24, 2025

Added this to extract values directly from the element vector

BryanCutler commented

View reviewed changes

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

+                      }
+                  }
+                  public void assignBlockFromStructVector(StructVector vector, Type type, BlockBuilder builder, int startIndex, int endIndex)

Contributor Author

BryanCutler Jan 24, 2025

Added this to extract values from all child vectors of the RowType

BryanCutler commented

View reviewed changes

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

+                      this.metadata = requireNonNull(metadata, "Metadata is null");
+                      this.splitManager = requireNonNull(splitManager, "splitManager is null");
+                      this.pageSourceProvider = requireNonNull(pageSourceProvider, "pageSourceProvider is null");
+                      this.connectorAllocator = requireNonNull(connectorAllocator, "connectorAllocator is null");

Contributor Author

BryanCutler Jan 24, 2025

Changed this to use a BufferAllocator instance. The connector will own this allocator and close it on shutdown. Flight clients created by the connector create a child allocator from this that is closed with the client.

BryanCutler commented

View reviewed changes

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

+                  ArrowErrorCode(int code, ErrorType type)
+                  {
+                      errorCode = new ErrorCode(code + 0x0510_0000, name(), type);

Contributor Author

BryanCutler Jan 24, 2025

changed error code range to be unique

BryanCutler commented

View reviewed changes

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

+                  public void configure(Binder binder)
+                  {
+                      configBinder(binder).bindConfig(ArrowFlightConfig.class);
+                      binder.bind(BufferAllocator.class).to(RootAllocator.class).in(Scopes.SINGLETON);

Contributor Author

BryanCutler Jan 24, 2025

Added this as the default RootAllocator instance, will be owned by the connector

BryanCutler commented

View reviewed changes

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

+                      this.columnHandles = requireNonNull(columnHandles, "columnHandles is null");
+                      requireNonNull(clientHandler, "clientHandler is null");
+                      this.arrowBlockBuilder = requireNonNull(arrowBlockBuilder, "arrowBlockBuilder is null");
+                      this.flightStreamAndClient = clientHandler.getFlightStream(split, connectorSession);

Contributor Author

BryanCutler Jan 24, 2025

Cleaned up the stream usage and closing here

BryanCutler commented

View reviewed changes

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

+                  @Override
+                  public Page getNextPage()
+                  {
+                      logger.debug("Reading next Arrow record batch");

Contributor Author

BryanCutler Jan 24, 2025

Added some debug logs

BryanCutler commented

View reviewed changes

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

+                      implements ConnectorTableHandle
+              {
+                  private final String schema;
+                  private final String table;

Contributor Author

BryanCutler Jan 24, 2025

There was previously discussion about making these optional #23032 (review)

BryanCutler commented

View reviewed changes

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

+              import static java.nio.file.Files.newInputStream;
+              import static java.util.Objects.requireNonNull;
+              public abstract class BaseArrowFlightClientHandler

Contributor Author

BryanCutler Jan 24, 2025

Renamed this to better indicate that it is a wrapper around the client and not an instance of the client, was confusing

BryanCutler commented

View reviewed changes

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

+                  public BaseArrowFlightClientHandler(BufferAllocator allocator, ArrowFlightConfig config)
+                  {
+                      this.allocator = requireNonNull(allocator, "allocator is null");

Contributor Author

BryanCutler Jan 24, 2025

Changed allocator usage here, it no longer creates a root allocator for each client instance. Instead an allocator is required in the constructor

BryanCutler commented

View reviewed changes

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

+                      }
+                  }
+                  protected ClientClosingFlightStream getFlightStream(ArrowSplit split, ConnectorSession connectorSession)

Contributor Author

BryanCutler Jan 24, 2025

Removed the Ticket parameter. Since it is in the ArrowSplit already, we can just deserialize once

BryanCutler force-pushed the arrow-connector-java-23032 branch from b530fd2 to 2d81e93 Compare

January 25, 2025 01:02

elbinpallimalilibm reviewed

View reviewed changes

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

+              public class TestingArrowFlightPlugin
+                      extends ArrowPlugin
+              {
+                  public TestingArrowFlightPlugin()

Contributor

elbinpallimalilibm Jan 25, 2025

this is part of the testing connector, should go into the testingConnector package

elbinpallimalilibm reviewed

View reviewed changes

...-flight/src/test/java/com/facebook/plugin/arrow/testingConnector/TestingH2DatabaseSetup.java

+              public class TestingH2DatabaseSetup
+              {
+                  private static final Logger logger = Logger.get(TestingH2DatabaseSetup.class);

Contributor

elbinpallimalilibm Jan 25, 2025

this is not a part of testing connector, can be moved to outside package

elbinpallimalilibm reviewed

View reviewed changes

...ow-flight/src/test/java/com/facebook/plugin/arrow/testingConnector/TestingArrowProducer.java

+              public class TestingArrowProducer
+                      implements FlightProducer
+              {
+                  private final BufferAllocator allocator;

Contributor

elbinpallimalilibm Jan 25, 2025

A connector implementation does not need this right? can be moved to outside package.

Contributor Author

BryanCutler Jan 27, 2025

I grouped everything needed by the connector into this package, including the server stuff. It seems like it would be a little clearer to move all testing server related classes into it's own package testingServer. I'll update the docs to explain this too.

elbinpallimalilibm reviewed

View reviewed changes

presto-docs/src/main/sphinx/connector/base-arrow-flight.rst Outdated

+              - ``ArrowBlockBuilder.java``
+                This class builds Presto blocks from Arrow vectors. Extend this class if needed and override ``getPrestoTypeFromArrowField`` method, if any customizations are needed for the conversion of Arrow vector to Presto type. A binding for this class should be created in the ``Module`` for the plugin.
+              A reference implementation of the presto-base-arrow-flight module in the test folder of this module uses a locally started Flight server to fetch data from an H2 database. The classes ``TestingArrowBlockBuilder``, ``TestingArrowFlightClientHandler``, ``TestingArrowFlightPlugin``, ``TestingArrowFlightRequest``, ``TestingArrowFlightResponse``, ``TestingArrowModule`` and ``TestingArrowQueryBuilder`` make up the reference implementation.

Contributor

elbinpallimalilibm Jan 25, 2025

Doc can refer to the testingConnector package to indicate the reference connector implementation

elbinpallimalilibm reviewed

View reviewed changes

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

+                  public void close()
+                  {
+                      if (flightStreamAndClient.hasRoot()) {
+                          completed = true;

Contributor

elbinpallimalilibm Jan 26, 2025 •

edited

Loading

Why set completed to true only when flight stream has root?

Contributor Author

BryanCutler Jan 27, 2025

This is how it was before, but looking into it further, it's not needed here. when flightStreamAndClient.next() returns false, that will set the completed flag, indicating that no more pages will be produced. I'll take this out.

tdcmeehan self-assigned this

BryanCutler force-pushed the arrow-connector-java-23032 branch 5 times, most recently from 59e58f4 to 9986850 Compare

January 28, 2025 03:43

tdcmeehan previously approved these changes

View reviewed changes

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

+                  public ArrowSplit(
+                          @JsonProperty("schemaName") @Nullable String schemaName,
+                          @JsonProperty("tableName") String tableName,
+                          @JsonProperty("flightEndpoint") byte[] flightEndpoint)

Contributor

tdcmeehan Jan 28, 2025

Let's call it flightEndpointBytes to indicate it's a serialized representation of the endpoint.

Contributor Author

BryanCutler Jan 29, 2025

Yeah, sounds better

...arrow-flight/src/test/java/com/facebook/plugin/arrow/testingServer/TestingArrowProducer.java Outdated

+                          }
+                          logger.info("Executing query: " + query);
+                          query = query.toUpperCase(); // Optionally, to maintain consistency

Contributor

tdcmeehan Jan 28, 2025

This comment doesn't appear to mean anything?

Contributor Author

BryanCutler Jan 29, 2025

I'm not really sure. I'll remove the comment and keep the uppercase conversion, in case that is preferred by the h2 driver.


          Adding Arrow Flight connector template

79c00d5

This adds the Arrow Flight base module as a template to build connectors
that use a Flight service to transfer data with Presto in Arrow format.
A concrete Arrow Flight connector implementation will extend the
BaseArrowFlightClientHandler that can connect to a Flight service that
handles the required RPC calls from the client handler. An example
connector implementation is provided in the unit testing.

RFC: https://github.com/prestodb/rfcs/blob/main/RFC-0004-arrow-flight-connector.md

Co-authored-by: sai bhaskar reddy <sai.bhaskar.reddy.sabbasani1@ibm.com>
Co-authored-by: SthuthiGhosh9400 <Sthuthi.Ghosh@ibm.com>
Co-authored-by: lithinwxd <Lithin.Purushothaman@ibm.com>
Co-authored-by: Steve Burnett <burnett@pobox.com>
Co-authored-by: elbinpallimalilibm <elbin.pallimalil@ibm.com>
Co-authored-by: Tim Meehan <tim@timdmeehan.com>

BryanCutler dismissed tdcmeehan’s stale review via

79c00d5

January 29, 2025 21:01

BryanCutler force-pushed the arrow-connector-java-23032 branch from 9986850 to 79c00d5 Compare

January 29, 2025 21:01

tdcmeehan approved these changes

View reviewed changes

BryanCutler merged commit 4464129 into prestodb:master

55 checks passed

BryanCutler deleted the arrow-connector-java-23032 branch

January 29, 2025 22:05

shangm2 pushed a commit to shangm2/presto that referenced this pull request


          Adding Arrow Flight connector template (prestodb#24427)

43a8a67

This adds the Arrow Flight base module as a template to build connectors
that use a Flight service to transfer data with Presto in Arrow format.
A concrete Arrow Flight connector implementation will extend the
BaseArrowFlightClientHandler that can connect to a Flight service that
handles the required RPC calls from the client handler. An example
connector implementation is provided in the unit testing.

RFC: https://github.com/prestodb/rfcs/blob/main/RFC-0004-arrow-flight-connector.md

Co-authored-by: sai bhaskar reddy <sai.bhaskar.reddy.sabbasani1@ibm.com>
Co-authored-by: SthuthiGhosh9400 <Sthuthi.Ghosh@ibm.com>
Co-authored-by: lithinwxd <Lithin.Purushothaman@ibm.com>
Co-authored-by: Steve Burnett <burnett@pobox.com>
Co-authored-by: elbinpallimalilibm <elbin.pallimalil@ibm.com>
Co-authored-by: Tim Meehan <tim@timdmeehan.com>

BryanCutler mentioned this pull request

[native] Add Apache Arrow Flight Connector #24504

Merged

6 tasks

steveburnett mentioned this pull request

Document the base arrow module #23212

Closed

6 tasks

BryanCutler mentioned this pull request

Apache Arrow Connector #12201

Closed

This was referenced Mar 10, 2025

Add release notes for 0.292 unix280/presto#5

Closed

Add release notes for 0.292 unix280/presto#6

Closed

jp-sivaprasad pushed a commit to jp-sivaprasad/presto that referenced this pull request


          Adding Arrow Flight connector template (prestodb#24427)

519e139

This adds the Arrow Flight base module as a template to build connectors
that use a Flight service to transfer data with Presto in Arrow format.
A concrete Arrow Flight connector implementation will extend the
BaseArrowFlightClientHandler that can connect to a Flight service that
handles the required RPC calls from the client handler. An example
connector implementation is provided in the unit testing.

RFC: https://github.com/prestodb/rfcs/blob/main/RFC-0004-arrow-flight-connector.md

Co-authored-by: sai bhaskar reddy <sai.bhaskar.reddy.sabbasani1@ibm.com>
Co-authored-by: SthuthiGhosh9400 <Sthuthi.Ghosh@ibm.com>
Co-authored-by: lithinwxd <Lithin.Purushothaman@ibm.com>
Co-authored-by: Steve Burnett <burnett@pobox.com>
Co-authored-by: elbinpallimalilibm <elbin.pallimalil@ibm.com>
Co-authored-by: Tim Meehan <tim@timdmeehan.com>

prestodb-ci mentioned this pull request

Add release notes for 0.292 #24825

Merged

30 tasks

unidevel mentioned this pull request

Add release notes for 0.292 unix280/presto#23

Closed

30 tasks

This was referenced May 6, 2025

Add release notes for 0.292 unix280/presto#28

Closed

Add release notes for 0.292 unix280/presto#29

Closed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

tdcmeehan tdcmeehan approved these changes

steveburnett steveburnett left review comments

elharo Awaiting requested review from elharo elharo is a code owner

presto-oss Awaiting requested review from presto-oss presto-oss is a code owner automatically assigned from prestodb/committers

aditi-pandit Awaiting requested review from aditi-pandit aditi-pandit was automatically assigned from prestodb/ibm-reviewers

psnv03 Awaiting requested review from psnv03 psnv03 was automatically assigned from prestodb/ibm-reviewers

+1 more reviewer

elbinpallimalilibm elbinpallimalilibm left review comments

Labels