diff --git a/lance-spark-base_2.12/src/main/java/com/lancedb/lance/spark/write/LanceArrowWriter.java b/lance-spark-base_2.12/src/main/java/com/lancedb/lance/spark/write/LanceArrowWriter.java
index 94d675a..00fde77 100644
--- a/lance-spark-base_2.12/src/main/java/com/lancedb/lance/spark/write/LanceArrowWriter.java
+++ b/lance-spark-base_2.12/src/main/java/com/lancedb/lance/spark/write/LanceArrowWriter.java
@@ -34,7 +34,7 @@ public class LanceArrowWriter extends ArrowReader {
   private final int batchSize;
 
   @GuardedBy("monitor")
-  private volatile boolean finished;
+  private volatile boolean finished = false;
 
   private final AtomicLong totalBytesRead = new AtomicLong();
   private com.lancedb.lance.spark.arrow.LanceArrowWriter arrowWriter = null;
@@ -70,8 +70,8 @@ void write(InternalRow row) {
   }
 
   void setFinished() {
-    loadToken.release();
     finished = true;
+    loadToken.release();
   }
 
   @Override
diff --git a/lance-spark-base_2.12/src/test/java/com/lancedb/lance/spark/write/LanceArrowWriterTest.java b/lance-spark-base_2.12/src/test/java/com/lancedb/lance/spark/write/LanceArrowWriterTest.java
index 019f6e1..2bea1dc 100644
--- a/lance-spark-base_2.12/src/test/java/com/lancedb/lance/spark/write/LanceArrowWriterTest.java
+++ b/lance-spark-base_2.12/src/test/java/com/lancedb/lance/spark/write/LanceArrowWriterTest.java
@@ -26,7 +26,6 @@
 import org.apache.spark.sql.types.DataTypes;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
 import java.util.Collections;
@@ -35,8 +34,6 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
-// TODO: Flaky in CI, needs to fix
-@Disabled
 public class LanceArrowWriterTest {
   @Test
   public void test() throws Exception {