Skip to content

Commit

Permalink
update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@vkorukanti
vkorukanti committed Jun 21, 2023
1 parent 7a93a71 commit 01acec2
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 16 deletions.
18 changes: 14 additions & 4 deletions kernel/kernel-default/src/main/java/io/delta/kernel/parquet/ParquetConverters.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -189,8 +189,9 @@ public void end()
public ColumnarBatch getDataAsColumnarBatch(int batchSize)
{
ColumnVector[] memberVectors = collectMemberVectors(batchSize);
this.currentRowIndex = 0;
return new DefaultColumnarBatch(batchSize, readSchema, memberVectors);
ColumnarBatch batch = new DefaultColumnarBatch(batchSize, readSchema, memberVectors);
resetWorkingState();
return batch;
}

@Override
Expand All @@ -214,12 +215,14 @@ public boolean moveToNextRow()
public ColumnVector getDataColumnVector(int batchSize)
{
ColumnVector[] memberVectors = collectMemberVectors(batchSize);
return new DefaultStructVector(
ColumnVector vector = new DefaultStructVector(
batchSize,
readSchema,
Optional.of(nullability),
memberVectors
);
resetWorkingState();
return vector;
}

private ColumnVector[] collectMemberVectors(int batchSize)
Expand All @@ -237,6 +240,12 @@ public void resizeIfNeeded()
this.nullability = Arrays.copyOf(this.nullability, newSize);
}
}

private void resetWorkingState()
{
this.currentRowIndex = 0;
this.nullability = initNullabilityVector(this.nullability.length);
}
}

public abstract static class BasePrimitiveColumnConverter
Expand Down Expand Up @@ -611,7 +620,8 @@ public void resizeIfNeeded()
}
}

static boolean[] initNullabilityVector(int size) {
static boolean[] initNullabilityVector(int size)
{
boolean[] nullability = new boolean[size];
// Initialize all values as null. As Parquet calls this converter only for non-null
// values, make the corresponding value to false.
Expand Down
50 changes: 38 additions & 12 deletions kernel/kernel-default/src/test/java/io/delta/kernel/parquet/TestParquetBatchReader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,10 @@
import java.util.List;
import java.util.Optional;
import org.apache.hadoop.conf.Configuration;
import org.junit.Assert;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import org.junit.Test;

Expand Down Expand Up @@ -184,7 +185,8 @@ private static void verifyRowFromAllTypesFile(
Boolean expValue = (rowId % 87 != 0) ? rowId % 2 == 0 : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertEquals(expValue.booleanValue(), vector.getBoolean(batchWithIdx._2));
}
break;
Expand All @@ -193,7 +195,8 @@ private static void verifyRowFromAllTypesFile(
Byte expValue = (rowId % 72 != 0) ? (byte) rowId : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertEquals(expValue.byteValue(), vector.getByte(batchWithIdx._2));
}
break;
Expand All @@ -202,7 +205,8 @@ private static void verifyRowFromAllTypesFile(
Short expValue = (rowId % 56 != 0) ? (short) rowId : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertEquals(expValue.shortValue(), vector.getShort(batchWithIdx._2));
}
break;
Expand All @@ -212,7 +216,8 @@ private static void verifyRowFromAllTypesFile(
new Date(rowId * 20000000L).toLocalDate() : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
long numDaysSinceEpoch = ChronoUnit.DAYS.between(EPOCH, expValue);
assertEquals(numDaysSinceEpoch, vector.getInt(batchWithIdx._2));
}
Expand All @@ -222,7 +227,8 @@ private static void verifyRowFromAllTypesFile(
Integer expValue = (rowId % 23 != 0) ? rowId : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertEquals(expValue.intValue(), vector.getInt(batchWithIdx._2));
}
break;
Expand All @@ -231,7 +237,8 @@ private static void verifyRowFromAllTypesFile(
Long expValue = (rowId % 25 != 0) ? rowId + 1L : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertEquals(expValue.longValue(), vector.getLong(batchWithIdx._2));
}
break;
Expand All @@ -240,7 +247,8 @@ private static void verifyRowFromAllTypesFile(
Float expValue = (rowId % 28 != 0) ? (rowId * 0.234f) : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertEquals(expValue.floatValue(), vector.getFloat(batchWithIdx._2), 0.02);
}
break;
Expand All @@ -249,16 +257,19 @@ private static void verifyRowFromAllTypesFile(
Double expValue = (rowId % 54 != 0) ? (rowId * 234234.23d) : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
assertEquals(expValue.doubleValue(), vector.getDouble(batchWithIdx._2), 0.02);
}
else {
assertEquals(expValue.doubleValue(), vector.getDouble(batchWithIdx._2),
0.02);
}
break;
}
case "stringtype": {
String expValue = (rowId % 57 != 0) ? Integer.toString(rowId) : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertArrayEquals(expValue.getBytes(), vector.getBinary(batchWithIdx._2));
}
break;
Expand All @@ -267,7 +278,8 @@ private static void verifyRowFromAllTypesFile(
byte[] expValue = (rowId % 57 != 0) ? Integer.toString(rowId).getBytes() : null;
if (expValue == null) {
assertTrue(vector.isNullAt(batchWithIdx._2));
} else {
}
else {
assertArrayEquals(expValue, vector.getBinary(batchWithIdx._2));
}
break;
Expand All @@ -280,6 +292,20 @@ private static void verifyRowFromAllTypesFile(
}
case "nested_struct": {
Row struct = vector.getStruct(batchWithIdx._2);
assertFalse(vector.isNullAt(batchWithIdx._2));
String aaVal = struct.getString(0);
assertEquals(Integer.toString(rowId), aaVal);

boolean expAcValNull = rowId % 23 == 0;
Row acVal = struct.getStruct(1);
if (expAcValNull) {
assertTrue(struct.isNullAt(1));
assertNull(acVal);
}
else {
int actAcaVal = acVal.getInt(0);
assertEquals(rowId, actAcaVal);
}
break;
}
case "array_of_prims": {
Expand Down

0 comments on commit 01acec2

Please sign in to comment.