ADH-5242

- fixed tests for decimal type with exceeding precision - fixed unloading data bug
arenadata · Dec 12, 2024 · 3470969 · 3470969
1 parent 5e90ba3
commit 3470969
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 139 deletions.
diff --git a/.../src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistResource.java b/.../src/main/java/io/trino/plugin/adb/connector/protocol/gpfdist/server/GpfdistResource.java
@@ -199,7 +199,7 @@ private void processInitialRequest(AsyncResponse asyncResponse, ReadContext read
     {
         InputDataProcessor dataProcessor = inputDataProcessorFactory.create(readContext.getRowDecoder(),
                 readContext.getRowProcessingService());
-        readContext.getSegmentDataProcessors().putIfAbsent(request.getSegmentId(),
+        readContext.getRowProcessingService().getSegmentDataProcessors().putIfAbsent(request.getSegmentId(),
                 new GpfdistSegmentRequestProcessor(request.getSegmentId(), dataProcessor));
         asyncResponse.resume(Response.ok()
                 .header(X_GP_PROTO, request.getGpProtocol())
@@ -238,7 +238,7 @@ private void processTearDownRequest(AsyncResponse asyncResponse, ReadContext rea
 
     private GpfdistSegmentRequestProcessor getSegmentProcessor(ReadContext readContext, Integer segmentId)
     {
-        return Optional.ofNullable(readContext.getSegmentDataProcessors().get(segmentId))
+        return Optional.ofNullable(readContext.getRowProcessingService().getSegmentDataProcessors().get(segmentId))
                 .orElseThrow(() -> new IllegalStateException(
                         "Failed to get segment request processor by segmentId: " + segmentId));
     }

diff --git a/...main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/RowProcessingService.java b/...main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/RowProcessingService.java
@@ -25,5 +25,7 @@ public interface RowProcessingService
 
     void stop();
 
+    void stopExceptionally(Throwable error);
+
     void clear();
 }
diff --git a/.../main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContext.java b/.../main/java/io/trino/plugin/adb/connector/protocol/gpfdist/unload/context/ReadContext.java
@@ -19,9 +19,7 @@
 import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.ContextId;
 import io.trino.plugin.adb.connector.protocol.gpfdist.metadata.GpfdistUnloadMetadata;
 import io.trino.plugin.adb.connector.protocol.gpfdist.unload.process.GpfdistBufferedRowProcessingService;
-import io.trino.plugin.adb.connector.protocol.gpfdist.unload.process.GpfdistSegmentRequestProcessor;
 
-import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
 
 public class ReadContext
@@ -63,11 +61,6 @@ public GpfdistUnloadMetadata getMetadata()
         return metadata;
     }
 
-    public Map<Integer, GpfdistSegmentRequestProcessor> getSegmentDataProcessors()
-    {
-        return rowProcessingService.getSegmentDataProcessors();
-    }
-
     public RowDecoder<String> getRowDecoder()
     {
         return rowDecoder;

diff --git a/...in/adb/connector/protocol/gpfdist/unload/process/GpfdistBufferedRowProcessingService.java b/...in/adb/connector/protocol/gpfdist/unload/process/GpfdistBufferedRowProcessingService.java
@@ -16,14 +16,14 @@
 import io.trino.plugin.adb.AdbPluginConfig;
 import io.trino.plugin.adb.connector.protocol.gpfdist.ConnectorRow;
 import io.trino.plugin.adb.connector.protocol.gpfdist.unload.RowProcessingService;
-import io.trino.plugin.adb.connector.protocol.gpfdist.unload.SegmentRequestStatus;
 
 import java.util.LinkedList;
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
@@ -34,6 +34,7 @@ public class GpfdistBufferedRowProcessingService
     private final Queue<ConnectorRow> rowsQueue = new LinkedList<>();
     private final AtomicLong usedBufferedRowsMemory = new AtomicLong(0);
     private final AtomicBoolean isStopped = new AtomicBoolean(false);
+    private final AtomicReference<Throwable> processingError = new AtomicReference<>();
     private final Map<Integer, GpfdistSegmentRequestProcessor> segmentDataProcessors = new ConcurrentHashMap<>();
     private final Condition isFullCondition;
     private final Condition isReadyForTransferCondition;
@@ -109,44 +110,48 @@ public void stop()
     {
         lock.lock();
         try {
-            isStopped.set(true);
-            isReadyForTransferCondition.signalAll();
+            stopInternal();
         }
         finally {
             lock.unlock();
         }
     }
 
     @Override
-    public void clear()
+    public void stopExceptionally(Throwable error)
     {
         lock.lock();
         try {
-            segmentDataProcessors.clear();
-            rowsQueue.clear();
+            stopInternal();
+            processingError.set(error);
         }
         finally {
             lock.unlock();
         }
     }
 
-    private boolean isNotReady()
+    private void stopInternal()
     {
-        return isDataTransferNotInitialized() || isNotAllDataProcessed();
+        isStopped.set(true);
+        isReadyForTransferCondition.signalAll();
     }
 
-    private boolean isDataTransferNotInitialized()
+    @Override
+    public void clear()
     {
-        //checked that insert adb query is not finished and there are no segment processors yet
-        return !isStopped.get() && segmentDataProcessors.isEmpty();
+        lock.lock();
+        try {
+            segmentDataProcessors.clear();
+            rowsQueue.clear();
+        }
+        finally {
+            lock.unlock();
+        }
     }
 
-    private boolean isNotAllDataProcessed()
+    private boolean isNotReady()
     {
-        //check that there are no rows in queue and segment processors are not finished yet
-        return isEmpty()
-                && segmentDataProcessors.values().stream()
-                .anyMatch(req -> req.getStatus() != SegmentRequestStatus.FINISHED);
+        return !isStopped.get() && (segmentDataProcessors.isEmpty() || isEmpty());
     }
 
     private boolean isFull()
@@ -168,4 +173,9 @@ public Map<Integer, GpfdistSegmentRequestProcessor> getSegmentDataProcessors()
     {
         return segmentDataProcessors;
     }
+
+    public AtomicReference<Throwable> getProcessingError()
+    {
+        return processingError;
+    }
 }
diff --git a/...va/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordCursor.java b/...va/io/trino/plugin/adb/connector/protocol/gpfdist/unload/process/GpfdistRecordCursor.java
@@ -75,8 +75,8 @@ public boolean advanceNextPosition()
                 }
                 currentRow = rowProcessingService.take();
 
-                if (queryExecutionException != null) {
-                    throw queryExecutionException;
+                if (rowProcessingService.getProcessingError().get() != null) {
+                    throw rowProcessingService.getProcessingError().get();
                 }
 
                 if (currentRow != null) {
@@ -88,8 +88,7 @@ public boolean advanceNextPosition()
                 else {
                     log.info("Data processing is finished. Unloaded rows %s. Processing result: %s",
                             unloadedRows,
-                            readContext.getSegmentDataProcessors().values());
-                    closeCtx();
+                            readContext.getRowProcessingService().getSegmentDataProcessors().values());
                     return false;
                 }
             }
@@ -99,7 +98,7 @@ public boolean advanceNextPosition()
                 }
                 log.warn("Processed rows to target: %d. Unloading from adb result: %s",
                         unloadedRows,
-                        readContext.getSegmentDataProcessors().values());
+                        readContext.getRowProcessingService().getSegmentDataProcessors().values());
                 closeCtx();
                 throw new TrinoException(StandardErrorCode.GENERIC_INTERNAL_ERROR,
                         "Failed to unload data from adb: " + e.getMessage(), e);
@@ -113,11 +112,14 @@ public boolean advanceNextPosition()
     private CompletableFuture<Void> executeTransferDataQuery()
     {
         return dataTransferQueryExecutor.execute().whenComplete((_, error) -> {
-            readContext.getRowProcessingService().stop();
+            log.debug("Data transfer query finished for contextId: %s", readContext.getId());
             if (error != null) {
-                queryExecutionException = error;
+                readContext.getRowProcessingService().stopExceptionally(error);
                 throw new CompletionException(error);
             }
+            else {
+                readContext.getRowProcessingService().stop();
+            }
         });
     }