Skip to content

Commit 2477cdc

Browse files
committed
fix: set finish before notifying load in LanceArrowWriter setFinished method
1 parent 2030a4e commit 2477cdc

File tree

2 files changed

+2
-5
lines changed

2 files changed

+2
-5
lines changed

lance-spark-base_2.12/src/main/java/com/lancedb/lance/spark/write/LanceArrowWriter.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public class LanceArrowWriter extends ArrowReader {
3434
private final int batchSize;
3535

3636
@GuardedBy("monitor")
37-
private volatile boolean finished;
37+
private volatile boolean finished = false;
3838

3939
private final AtomicLong totalBytesRead = new AtomicLong();
4040
private com.lancedb.lance.spark.arrow.LanceArrowWriter arrowWriter = null;
@@ -70,8 +70,8 @@ void write(InternalRow row) {
7070
}
7171

7272
void setFinished() {
73-
loadToken.release();
7473
finished = true;
74+
loadToken.release();
7575
}
7676

7777
@Override

lance-spark-base_2.12/src/test/java/com/lancedb/lance/spark/write/LanceArrowWriterTest.java

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import org.apache.spark.sql.types.DataTypes;
2727
import org.apache.spark.sql.types.StructField;
2828
import org.apache.spark.sql.types.StructType;
29-
import org.junit.jupiter.api.Disabled;
3029
import org.junit.jupiter.api.Test;
3130

3231
import java.util.Collections;
@@ -35,8 +34,6 @@
3534

3635
import static org.junit.jupiter.api.Assertions.assertEquals;
3736

38-
// TODO: Flaky in CI, needs to fix
39-
@Disabled
4037
public class LanceArrowWriterTest {
4138
@Test
4239
public void test() throws Exception {

0 commit comments

Comments
 (0)