Add the base FlightProducer for getStream API #17446
Conversation
github-actions
bot
added
enhancement
| public class ArrowFlightServerIT extends OpenSearchIntegTestCase { | ||
|
|
||
| private FlightClientManager flightClientManager; | ||
|
|
||
| @BeforeClass | ||
| public static void setupFeatureFlags() { | ||
| FeatureFlagSetter.set(FeatureFlags.ARROW_STREAMS_SETTING.getKey()); |
There was a problem hiding this comment.
for OpenSearchIntegTestCase you can override featureFlagSettings() to provide the flags as node level settings that won't mess with system properties like FeatureFlagSetter:
example:
@Override
protected Settings featureFlagSettings() {
return Settings.builder().put(super.featureFlagSettings()).put(FeatureFlags. ARROW_STREAMS_SETTING, Boolean.TRUE).build();
}
There was a problem hiding this comment.
thanks, let me try it
There was a problem hiding this comment.
seems like this way doesn't work for integ tests, also it seems like FeatureFlagSetter is purpose built for it.
| infoBuilder = FlightInfo.builder( | ||
| streamProducerHolder.get().getRoot().getSchema(), | ||
| descriptor, | ||
| Collections.singletonList(endpoint) |
There was a problem hiding this comment.
Curious what your thoughts on providing additional endpoints via the StreamProducer to consolidate to a single ticket at a coordinator? In a poc I've been working on with DataFusion I am using this to specify partitions (each shard ticket) from FlightInfo and only hand a single ticket to the engine and then it calls getStream for each partition. I could see this also helping with join case where we get a ticket per index or further consolidating to a single ticket to pass back to a client.
There was a problem hiding this comment.
@mch2 taking a look
There was a problem hiding this comment.
This requirement looks valid. So far, StreamTicket represents a unit of work, thus each ticket is associated with single StreamProducer. However, this unit of work could be partitioned, I would like to solve it by creating abstraction at StreamTicket itself -
we can create something like PartitionedStreamTicket, which will internally contain list of tickets.
Its StreamProducer implementation in java can have a parallel execution of streams when getStream() is called on it (you may not use it though in data fusion).
For getFlightInfo() it can just return List of endpoints, as your POC code.
Let me know what do you think? created - #17493
There was a problem hiding this comment.
Sounds reasonable, but you'd need a worker/coordinator to create that single ticket after receiving each partition's (data node) tickets post query phase. At least in the data node case you wouldn't know those tickets beforehand.
| ); | ||
| return; | ||
| } | ||
| StreamProducer<VectorSchemaRoot, BufferAllocator> proxyProvider = new ProxyStreamProducer( |
There was a problem hiding this comment.
am curious when this would happen given getStream is typically invoked after a getFlightInfo call that vends the endpoint? Or is the intent to skip that step and allow hitting any node with a ticket and stream a reply? Would that not add latency & load on an additional data node thats vending through the proxy that we could otherwise avoid?
There was a problem hiding this comment.
@mch2 for communication within cluster nodes, i think you're right as we should not make use of ProxyStreamProducer. This will be useful when clients directly want to connect via load balancer and do not have access to the nodes directly, so any node can behave as a coordinator and then it should know where to route.
For cases, where data nodes are directly accessible, this intelligence should be in the client to call getFlightInfo and send the request directly to the node holding stream.
rishabhmaurya
requested review from
anasalkouz,
andrross,
ashking94,
bugmakerrrrrr,
Bukhtawar,
CEHENKLE,
cwperks,
dblock,
dbwiddis,
gbbafna,
jainankitk,
kotwanikunal,
linuxpi,
msfroh,
nknize,
owaiskazi19,
reta,
Rishikesh1159,
sachinpkale,
saratvemulapalli,
shwetathareja and
sohami
as code owners
rishabhmaurya
force-pushed
the
flight-base-impl
branch
from
7738286 to
c711cf6
Compare
|
❌ Gradle check result for c711cf6: FAILURE Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change? |
Signed-off-by: Rishabh Maurya <[email protected]>
Signed-off-by: Rishabh Maurya <[email protected]>
Signed-off-by: Rishabh Maurya <[email protected]>
rishabhmaurya
force-pushed
the
flight-base-impl
branch
from
c711cf6 to
2532e0f
Compare
|
❌ Gradle check result for 2532e0f: FAILURE Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change? |
Signed-off-by: Rishabh Maurya <[email protected]>
rishabhmaurya
force-pushed
the
flight-base-impl
branch
from
2532e0f to
44423f1
Compare
|
❌ Gradle check result for 44423f1: FAILURE Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change? |
| @@ -88,6 +88,7 @@ internalClusterTest { | |||
| systemProperty 'io.netty.noUnsafe', 'false' | |||
| systemProperty 'io.netty.tryUnsafe', 'true' | |||
| systemProperty 'io.netty.tryReflectionSetAccessible', 'true' | |||
| jvmArgs += ["--add-opens", "java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED"] | |||
There was a problem hiding this comment.
I think this should go to jvm.options as well:
21-:--add-opens java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED
| @@ -99,7 +100,19 @@ public FlightClientManager( | |||
| * @return An OpenSearchFlightClient instance for the specified node | |||
| */ | |||
| public Optional<FlightClient> getFlightClient(String nodeId) { | |||
| return Optional.ofNullable(flightClients.get(nodeId)); | |||
| ClientHolder clientHolder = flightClients.getOrDefault(nodeId, null); | |||
There was a problem hiding this comment.
You will always get null when key is not in the map:
| ClientHolder clientHolder = flightClients.getOrDefault(nodeId, null); | |
| ClientHolder clientHolder = flightClients.get(nodeId); |
| @@ -99,7 +100,19 @@ public FlightClientManager( | |||
| * @return An OpenSearchFlightClient instance for the specified node | |||
| */ | |||
| public Optional<FlightClient> getFlightClient(String nodeId) { | |||
| return Optional.ofNullable(flightClients.get(nodeId)); | |||
| ClientHolder clientHolder = flightClients.getOrDefault(nodeId, null); | |||
| return clientHolder == null ? Optional.empty() : Optional.ofNullable(clientHolder.flightClient); | |||
There was a problem hiding this comment.
I think the flightClient should never be null, right?
| return clientHolder == null ? Optional.empty() : Optional.ofNullable(clientHolder.flightClient); | |
| return clientHolder == null ? Optional.empty() : Optional.of(clientHolder.flightClient); |
| * @param nodeId The ID of the node for which to retrieve the location | ||
| * @return The Location of the Flight client for the specified node | ||
| */ | ||
| public Location getFlightClientLocation(String nodeId) { |
There was a problem hiding this comment.
Please, use Optional here as well:
| public Location getFlightClientLocation(String nodeId) { | |
| public Optional<Location> getFlightClientLocation(String nodeId) { |
| * @return The Location of the Flight client for the specified node | ||
| */ | ||
| public Location getFlightClientLocation(String nodeId) { | ||
| ClientHolder clientHolder = flightClients.getOrDefault(nodeId, null); |
There was a problem hiding this comment.
| ClientHolder clientHolder = flightClients.getOrDefault(nodeId, null); | |
| ClientHolder clientHolder = flightClients.get(nodeId); |
| future.completeExceptionally(new IllegalStateException("No Flight info received for node: [" + nodeId + "]")); | ||
| } | ||
| } | ||
| client.execute(NodesFlightInfoAction.INSTANCE, request, new ActionListener<>() { |
There was a problem hiding this comment.
Why we removed try from here?
| @@ -184,13 +200,16 @@ private DiscoveryNode getNodeFromClusterState(String nodeId) { | |||
| */ | |||
| @Override | |||
| public void close() throws Exception { | |||
| for (FlightClient flightClient : flightClients.values()) { | |||
| flightClient.close(); | |||
| closed = true; | |||
| @@ -229,7 +248,7 @@ private Set<String> getCurrentClusterNodes() { | |||
| } | |||
|
|
|||
| @VisibleForTesting | |||
| Map<String, FlightClient> getFlightClients() { | |||
| Map<String, ClientHolder> getFlightClients() { | |||
There was a problem hiding this comment.
You already have getClients that does the same?
| @@ -165,6 +166,7 @@ protected void doClose() { | |||
| } | |||
|
|
|||
| private void initializeStreamManager(FlightClientManager clientManager) { | |||
| streamManager = null; | |||
| streamManager = new FlightStreamManager(() -> allocator); | |||
There was a problem hiding this comment.
Please change type of the field: StreamManager streamManager -> FlightStreamManager streamManager, no need to typecast all the time since the streamManager is created here
Description
With recently added Arrow SPIs #16691 and flight server support #16962, this is a follow up change which introduces -
getStream()andgetFlightInfo()APIs. It also adds support for defaultStreamManager,FlightProducer,FlightStreamReaderwhich these APIs make use of.StreamManagerPluginto get instance of defaultStreamManager.ProxyStreamProviderwhich acts as forward proxy for FlightStream. This is useful when stream is not present locally and needs to be fetched from a different node in the cluster.Related Issues
Resolves #17065
Check List
- [ ] API changes companion pull request created, if applicable.- [ ] Public documentation issue/PR created, if applicable.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.