github-actions[bot] commented on code in PR #41315:
URL: https://github.com/apache/beam/pull/41315#discussion_r4567117402
##########
sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java:
##########
@@ -180,10 +185,132 @@ static Enumerable<Object> toEnumerable(
return count(options, node, cancelFlag);
} else if (isLimitQuery(node)) {
return limitCollect(options, node, cancelFlag);
+ } else if (node.isBounded() == IsBounded.UNBOUNDED) {
+ return liveCollect(options, node, cancelFlag);
}
return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues(collectRows(options,
node, cancelFlag)));
}
+ /**
+ * Runs the pipeline of an unbounded query and hands over its rows while it
runs.
+ *
+ * <p>If the shell set a {@link LiveResult.Viewer}, it shows the rows and
the query returns none.
+ * Otherwise the query returns the rows as they arrive, until the pipeline
ends, the statement is
+ * cancelled, or the result set is closed.
+ */
+ private static Enumerable<Object> liveCollect(
+ PipelineOptions options, BeamRelNode node, AtomicBoolean cancelFlag) {
+ checkArgument(
+ options
+ .getRunner()
+ .getCanonicalName()
+ .equals("org.apache.beam.runners.direct.DirectRunner"),
+ "SELECT without INSERT is only supported in DirectRunner in SQL
Shell.");
+ Pipeline pipeline = Pipeline.create(options);
+ PCollection<Row> rows = BeamSqlRelUtils.toPCollection(pipeline, node);
+ rows.apply(ParDo.of(new LiveCollector()));
+ LiveResult.Viewer viewer = LiveResult.getViewer();
+ LiveResult result = LiveResult.start(pipeline, rows.getSchema(),
cancelFlag);
+ if (viewer == null) {
+ return new LiveRows(result);
+ }
+
+ RuntimeException failure;
+ try {
+ viewer.show(result);
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to show the rows of the query", e);
+ } finally {
+ // Read the failure first: errors after closing come from cancelling the
pipeline.
+ failure = result.getFailure();
+ result.close();
+ }
+ if (failure != null) {
+ throw failure;
+ }
+ return Linq4j.emptyEnumerable();
+ }
+
+ /** The rows of a {@link LiveResult} as they arrive. Closing them cancels
the pipeline. */
+ private static class LiveRows extends AbstractEnumerable<Object> {
+ private final LiveResult result;
+
+ LiveRows(LiveResult result) {
+ this.result = result;
+ }
+
+ @Override
Review Comment:
**🟠High**: `LiveRows` extends `AbstractEnumerable<Object>` (which
implements `Enumerable<Object>`). In Calcite's
`Linq4j.iterableEnumerator(Iterable)`, the first check is `if (iterable
instanceof Enumerable) { return ((Enumerable<T>) iterable).enumerator(); }`.
Passing `this` causes infinite recursion (`StackOverflowError`) whenever
`enumerator()` is called on `LiveRows`. Pass `this::iterator` (or extend
`AbstractEnumerable2<Object>`) so `Linq4j` wraps the `Iterator` instead of
recursing back into `enumerator()`.
```suggestion
return Linq4j.iterableEnumerator(this::iterator);
```
##########
sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/LiveResult.java:
##########
@@ -0,0 +1,228 @@
+/*
+ * 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.beam.sdk.extensions.sql.impl.rel;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.beam.runners.direct.DirectOptions;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.PipelineResult.State;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * The rows of an unbounded query, handed over while its pipeline runs.
+ *
+ * <p>{@link BeamEnumerableConverter} runs the pipeline on the DirectRunner.
Its last step adds each
+ * row and its {@link ValueKind} here, and a JDBC result set or the {@link
Viewer} takes them out.
+ * This only works on the DirectRunner, because it runs the pipeline in this
JVM.
+ */
+@Internal
+public final class LiveResult {
+ /** Shows the rows of unbounded queries while their pipelines run, in place
of JDBC results. */
+ public interface Viewer {
+ /**
+ * Shows the rows until the user is done or {@link
LiveResult#isCancelRequested()} turns true.
+ * It runs on the statement's thread. Once it returns, the pipeline is
cancelled and the
+ * statement returns no rows, or fails if the pipeline failed.
+ */
+ void show(LiveResult result) throws IOException;
+ }
+
+ /** A row and the kind of change it is. */
+ public static final class Change {
+ private final ValueKind kind;
+ private final Row row;
+
+ Change(ValueKind kind, Row row) {
+ this.kind = kind;
+ this.row = row;
+ }
+
+ public ValueKind getKind() {
+ return kind;
+ }
+
+ public Row getRow() {
+ return row;
+ }
+ }
+
+ private static final Logger LOG = LoggerFactory.getLogger(LiveResult.class);
+
+ /** How many rows wait to be taken out before the pipeline waits for room. */
+ private static final int CAPACITY = 10_000;
+
+ /** The live results of running pipelines, by pipeline options id. */
+ private static final Map<Long, LiveResult> RESULTS = new
ConcurrentHashMap<>();
+
+ private static volatile @Nullable Viewer viewer;
+
+ private final Schema schema;
+ private final AtomicBoolean cancelFlag;
+ private final long id;
+ private final LinkedBlockingDeque<Change> changes = new
LinkedBlockingDeque<>(CAPACITY);
+ private final CountDownLatch finished = new CountDownLatch(1);
+ private volatile @Nullable PipelineResult pipelineResult;
+ private volatile State state = State.RUNNING;
+ private volatile @Nullable RuntimeException failure;
+ private volatile boolean closed;
+
+ private LiveResult(Schema schema, AtomicBoolean cancelFlag, long id) {
+ this.schema = schema;
+ this.cancelFlag = cancelFlag;
+ this.id = id;
+ }
+
+ /**
+ * Sets the viewer for the rows of unbounded queries. The Beam SQL shell
sets one when it runs in
+ * a terminal. Without one, the rows are returned through JDBC as they
arrive.
+ */
+ public static void setViewer(@Nullable Viewer newViewer) {
+ viewer = newViewer;
+ }
+
+ static @Nullable Viewer getViewer() {
+ return viewer;
+ }
+
+ static @Nullable LiveResult forOptionsId(long optionsId) {
+ return RESULTS.get(optionsId);
+ }
+
+ /**
+ * Starts the pipeline without waiting for it, and watches it until it ends.
+ *
+ * @param cancelFlag set when the statement is cancelled, for example with
Ctrl+C in the shell
+ */
+ static LiveResult start(Pipeline pipeline, Schema schema, AtomicBoolean
cancelFlag) {
+ LiveResult result = new LiveResult(schema, cancelFlag,
pipeline.getOptions().getOptionsId());
+ RESULTS.put(result.id, result);
+ PipelineResult pipelineResult;
+ try {
+ pipeline.getOptions().as(DirectOptions.class).setBlockOnRun(false);
+ pipelineResult = pipeline.run();
+ } catch (RuntimeException e) {
+ RESULTS.remove(result.id);
+ throw e;
+ }
+ result.pipelineResult = pipelineResult;
+ Thread watcher = new Thread(() -> result.watch(pipelineResult),
"BeamSqlLiveResultWatcher");
+ watcher.setDaemon(true);
+ watcher.start();
+ return result;
+ }
+
+ private void watch(PipelineResult pipelineResult) {
+ @Nullable State end = null;
+ try {
+ end = pipelineResult.waitUntilFinish();
+ } catch (RuntimeException e) {
Review Comment:
**🟡 Medium**: `DirectRunner`'s
`ExecutorServiceParallelExecutor.waitUntilFinish()` rethrows `Error` subclasses
(e.g., `AssertionError`, `NoClassDefFoundError`) directly without wrapping them
in a `RuntimeException`. Because `watch()` only catches `RuntimeException`, an
`Error` escapes the `try` block, sets `state = State.UNKNOWN` and counts down
`finished` in `finally` while leaving `failure = null`. The consumer then sees
`isFinished() == true` and `getFailure() == null`, silently treating a crashed
pipeline as a normal end-of-stream. Additionally, if `waitUntilFinish()`
returns `State.FAILED` without throwing, `failure` remains `null`. Catch
`Throwable` (wrapping non-`RuntimeException` in a `RuntimeException`) and
ensure `failure` is populated if `end == State.FAILED`.
##########
sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamEnumerableConverter.java:
##########
@@ -180,10 +185,132 @@ static Enumerable<Object> toEnumerable(
return count(options, node, cancelFlag);
} else if (isLimitQuery(node)) {
return limitCollect(options, node, cancelFlag);
+ } else if (node.isBounded() == IsBounded.UNBOUNDED) {
+ return liveCollect(options, node, cancelFlag);
}
return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues(collectRows(options,
node, cancelFlag)));
}
+ /**
+ * Runs the pipeline of an unbounded query and hands over its rows while it
runs.
+ *
+ * <p>If the shell set a {@link LiveResult.Viewer}, it shows the rows and
the query returns none.
+ * Otherwise the query returns the rows as they arrive, until the pipeline
ends, the statement is
+ * cancelled, or the result set is closed.
+ */
+ private static Enumerable<Object> liveCollect(
+ PipelineOptions options, BeamRelNode node, AtomicBoolean cancelFlag) {
+ checkArgument(
+ options
+ .getRunner()
+ .getCanonicalName()
+ .equals("org.apache.beam.runners.direct.DirectRunner"),
+ "SELECT without INSERT is only supported in DirectRunner in SQL
Shell.");
+ Pipeline pipeline = Pipeline.create(options);
+ PCollection<Row> rows = BeamSqlRelUtils.toPCollection(pipeline, node);
+ rows.apply(ParDo.of(new LiveCollector()));
+ LiveResult.Viewer viewer = LiveResult.getViewer();
+ LiveResult result = LiveResult.start(pipeline, rows.getSchema(),
cancelFlag);
+ if (viewer == null) {
+ return new LiveRows(result);
+ }
+
+ RuntimeException failure;
+ try {
+ viewer.show(result);
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to show the rows of the query", e);
+ } finally {
+ // Read the failure first: errors after closing come from cancelling the
pipeline.
+ failure = result.getFailure();
+ result.close();
+ }
+ if (failure != null) {
+ throw failure;
+ }
+ return Linq4j.emptyEnumerable();
+ }
+
+ /** The rows of a {@link LiveResult} as they arrive. Closing them cancels
the pipeline. */
+ private static class LiveRows extends AbstractEnumerable<Object> {
+ private final LiveResult result;
+
+ LiveRows(LiveResult result) {
+ this.result = result;
+ }
+
+ @Override
+ public Enumerator<Object> enumerator() {
+ return Linq4j.iterableEnumerator(this);
+ }
+
+ /**
+ * JDBC reads the rows through this iterator. Unlike the default one, it
doesn't wait for the
+ * next row before it returns the current one, which would hold back the
newest row.
+ */
+ @Override
+ public Iterator<Object> iterator() {
+ return new LiveIterator();
+ }
+
+ private class LiveIterator implements Iterator<Object>, AutoCloseable {
+ private LiveResult.@Nullable Change next;
+
+ @Override
+ public boolean hasNext() {
+ // A cancelled statement ends the rows.
+ while (next == null && !result.isCancelRequested()) {
+ // No rows come after the pipeline ends, so then an empty poll means
there are no more.
+ boolean finished = result.isFinished();
Review Comment:
**🟢 Low**: When `finished` is already `true` and the buffer is empty,
`result.poll(100, TimeUnit.MILLISECONDS)` still blocks for 100ms even though no
further rows can be produced. Using a timeout of `0` when `finished` is `true`
avoids adding a 100ms delay at the end of every bounded-completion or cancelled
query.
```suggestion
next = result.poll(finished ? 0 : 100, TimeUnit.MILLISECONDS);
```
##########
sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/LiveResult.java:
##########
@@ -0,0 +1,228 @@
+/*
+ * 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.beam.sdk.extensions.sql.impl.rel;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.beam.runners.direct.DirectOptions;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.PipelineResult.State;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.ValueKind;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * The rows of an unbounded query, handed over while its pipeline runs.
+ *
+ * <p>{@link BeamEnumerableConverter} runs the pipeline on the DirectRunner.
Its last step adds each
+ * row and its {@link ValueKind} here, and a JDBC result set or the {@link
Viewer} takes them out.
+ * This only works on the DirectRunner, because it runs the pipeline in this
JVM.
+ */
+@Internal
+public final class LiveResult {
+ /** Shows the rows of unbounded queries while their pipelines run, in place
of JDBC results. */
+ public interface Viewer {
+ /**
+ * Shows the rows until the user is done or {@link
LiveResult#isCancelRequested()} turns true.
+ * It runs on the statement's thread. Once it returns, the pipeline is
cancelled and the
+ * statement returns no rows, or fails if the pipeline failed.
+ */
+ void show(LiveResult result) throws IOException;
+ }
+
+ /** A row and the kind of change it is. */
+ public static final class Change {
+ private final ValueKind kind;
+ private final Row row;
+
+ Change(ValueKind kind, Row row) {
+ this.kind = kind;
+ this.row = row;
+ }
+
+ public ValueKind getKind() {
+ return kind;
+ }
+
+ public Row getRow() {
+ return row;
+ }
+ }
+
+ private static final Logger LOG = LoggerFactory.getLogger(LiveResult.class);
+
+ /** How many rows wait to be taken out before the pipeline waits for room. */
+ private static final int CAPACITY = 10_000;
+
+ /** The live results of running pipelines, by pipeline options id. */
+ private static final Map<Long, LiveResult> RESULTS = new
ConcurrentHashMap<>();
+
+ private static volatile @Nullable Viewer viewer;
+
+ private final Schema schema;
+ private final AtomicBoolean cancelFlag;
+ private final long id;
+ private final LinkedBlockingDeque<Change> changes = new
LinkedBlockingDeque<>(CAPACITY);
+ private final CountDownLatch finished = new CountDownLatch(1);
+ private volatile @Nullable PipelineResult pipelineResult;
+ private volatile State state = State.RUNNING;
+ private volatile @Nullable RuntimeException failure;
+ private volatile boolean closed;
+
+ private LiveResult(Schema schema, AtomicBoolean cancelFlag, long id) {
+ this.schema = schema;
+ this.cancelFlag = cancelFlag;
+ this.id = id;
+ }
+
+ /**
+ * Sets the viewer for the rows of unbounded queries. The Beam SQL shell
sets one when it runs in
+ * a terminal. Without one, the rows are returned through JDBC as they
arrive.
+ */
+ public static void setViewer(@Nullable Viewer newViewer) {
+ viewer = newViewer;
+ }
+
+ static @Nullable Viewer getViewer() {
+ return viewer;
+ }
+
+ static @Nullable LiveResult forOptionsId(long optionsId) {
+ return RESULTS.get(optionsId);
+ }
+
+ /**
+ * Starts the pipeline without waiting for it, and watches it until it ends.
+ *
+ * @param cancelFlag set when the statement is cancelled, for example with
Ctrl+C in the shell
+ */
+ static LiveResult start(Pipeline pipeline, Schema schema, AtomicBoolean
cancelFlag) {
+ LiveResult result = new LiveResult(schema, cancelFlag,
pipeline.getOptions().getOptionsId());
+ RESULTS.put(result.id, result);
+ PipelineResult pipelineResult;
+ try {
+ pipeline.getOptions().as(DirectOptions.class).setBlockOnRun(false);
+ pipelineResult = pipeline.run();
+ } catch (RuntimeException e) {
+ RESULTS.remove(result.id);
+ throw e;
+ }
+ result.pipelineResult = pipelineResult;
+ Thread watcher = new Thread(() -> result.watch(pipelineResult),
"BeamSqlLiveResultWatcher");
+ watcher.setDaemon(true);
+ watcher.start();
+ return result;
+ }
+
+ private void watch(PipelineResult pipelineResult) {
+ @Nullable State end = null;
+ try {
+ end = pipelineResult.waitUntilFinish();
+ } catch (RuntimeException e) {
+ failure = e;
+ end = State.FAILED;
+ } finally {
+ state = end == null ? State.UNKNOWN : end;
+ finished.countDown();
+ }
+ }
+
+ /**
+ * Adds a row, waiting for room if the buffer is full. The pipeline calls
this from many threads.
+ */
+ void add(ValueKind kind, Row row) {
+ Change change = new Change(kind, row);
+ try {
+ while (!closed) {
Review Comment:
**🟡 Medium**: If `Statement.cancel()` is called while the caller is *not*
blocked inside `LiveIterator.hasNext()` (for example, between
`ResultSet.next()` calls), `AvaticaResultSet.next()` checks
`statement.cancelFlag.get()` before calling `cursor.next()` and throws
`SQLException("Statement canceled")` without ever entering
`LiveIterator.hasNext()`. In that case `LiveResult.close()` is not triggered by
`hasNext()`, and worker threads blocked in `add()` when the queue is full will
stay blocked until the `ResultSet` is explicitly closed. Checking
`!cancelFlag.get()` here (and/or monitoring `cancelFlag` in `watch`) avoids
blocking pipeline workers after cancellation is requested.
```suggestion
while (!closed && !cancelFlag.get()) {
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
