Skip to content

Commit ec42a3a

Browse files
walterddrRong Rong
andauthored
fix column data type store for data table (#8648)
* fix column data type store for data table * use double for bigdecimal gen Co-authored-by: Rong Rong <[email protected]>
1 parent 6192229 commit ec42a3a

File tree

2 files changed

+75
-4
lines changed

2 files changed

+75
-4
lines changed

pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,10 @@ public ColumnDataType getStoredType() {
277277
return LONG;
278278
case JSON:
279279
return STRING;
280+
case BOOLEAN_ARRAY:
281+
return INT_ARRAY;
282+
case TIMESTAMP_ARRAY:
283+
return LONG_ARRAY;
280284
default:
281285
return this;
282286
}

pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java

Lines changed: 71 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.io.ByteArrayInputStream;
2525
import java.io.DataInputStream;
2626
import java.io.IOException;
27+
import java.math.BigDecimal;
2728
import java.nio.ByteBuffer;
2829
import java.util.Arrays;
2930
import java.util.Map;
@@ -57,13 +58,19 @@ public class DataTableSerDeTest {
5758
private static final long[] LONGS = new long[NUM_ROWS];
5859
private static final float[] FLOATS = new float[NUM_ROWS];
5960
private static final double[] DOUBLES = new double[NUM_ROWS];
61+
private static final BigDecimal[] BIG_DECIMALS = new BigDecimal[NUM_ROWS];
62+
private static final int[] BOOLEANS = new int[NUM_ROWS];
63+
private static final long[] TIMESTAMPS = new long[NUM_ROWS];
6064
private static final String[] STRINGS = new String[NUM_ROWS];
65+
private static final String[] JSONS = new String[NUM_ROWS];
6166
private static final byte[][] BYTES = new byte[NUM_ROWS][];
6267
private static final Object[] OBJECTS = new Object[NUM_ROWS];
6368
private static final int[][] INT_ARRAYS = new int[NUM_ROWS][];
6469
private static final long[][] LONG_ARRAYS = new long[NUM_ROWS][];
6570
private static final float[][] FLOAT_ARRAYS = new float[NUM_ROWS][];
6671
private static final double[][] DOUBLE_ARRAYS = new double[NUM_ROWS][];
72+
private static final int[][] BOOLEAN_ARRAYS = new int[NUM_ROWS][];
73+
private static final long[][] TIMESTAMP_ARRAYS = new long[NUM_ROWS][];
6774
private static final String[][] STRING_ARRAYS = new String[NUM_ROWS][];
6875
private static final Map<String, String> EXPECTED_METADATA =
6976
ImmutableMap.<String, String>builder().put(MetadataKey.NUM_DOCS_SCANNED.getName(), String.valueOf(20L))
@@ -417,10 +424,26 @@ private void fillDataTableWithRandomData(DataTableBuilder dataTableBuilder,
417424
DOUBLES[rowId] = RANDOM.nextDouble();
418425
dataTableBuilder.setColumn(colId, DOUBLES[rowId]);
419426
break;
427+
case BIG_DECIMAL:
428+
BIG_DECIMALS[rowId] = BigDecimal.valueOf(RANDOM.nextDouble());
429+
dataTableBuilder.setColumn(colId, BIG_DECIMALS[rowId]);
430+
break;
431+
case TIMESTAMP:
432+
TIMESTAMPS[rowId] = RANDOM.nextLong();
433+
dataTableBuilder.setColumn(colId, TIMESTAMPS[rowId]);
434+
break;
435+
case BOOLEAN:
436+
BOOLEANS[rowId] = RANDOM.nextInt(2);
437+
dataTableBuilder.setColumn(colId, BOOLEANS[rowId]);
438+
break;
420439
case STRING:
421440
STRINGS[rowId] = RandomStringUtils.random(RANDOM.nextInt(20));
422441
dataTableBuilder.setColumn(colId, STRINGS[rowId]);
423442
break;
443+
case JSON:
444+
JSONS[rowId] = "{\"key\": \"" + RandomStringUtils.random(RANDOM.nextInt(20)) + "\"}";
445+
dataTableBuilder.setColumn(colId, JSONS[rowId]);
446+
break;
424447
case BYTES:
425448
BYTES[rowId] = RandomStringUtils.random(RANDOM.nextInt(20)).getBytes();
426449
dataTableBuilder.setColumn(colId, new ByteArray(BYTES[rowId]));
@@ -466,6 +489,27 @@ private void fillDataTableWithRandomData(DataTableBuilder dataTableBuilder,
466489
DOUBLE_ARRAYS[rowId] = doubleArray;
467490
dataTableBuilder.setColumn(colId, doubleArray);
468491
break;
492+
case BOOLEAN_ARRAY:
493+
length = RANDOM.nextInt(2);
494+
int[] booleanArray = new int[length];
495+
for (int i = 0; i < length; i++) {
496+
booleanArray[i] = RANDOM.nextInt();
497+
}
498+
BOOLEAN_ARRAYS[rowId] = booleanArray;
499+
dataTableBuilder.setColumn(colId, booleanArray);
500+
break;
501+
case TIMESTAMP_ARRAY:
502+
length = RANDOM.nextInt(20);
503+
long[] timestampArray = new long[length];
504+
for (int i = 0; i < length; i++) {
505+
timestampArray[i] = RANDOM.nextLong();
506+
}
507+
TIMESTAMP_ARRAYS[rowId] = timestampArray;
508+
dataTableBuilder.setColumn(colId, timestampArray);
509+
break;
510+
case BYTES_ARRAY:
511+
// TODO: add once implementation of datatable bytes array support is added
512+
break;
469513
case STRING_ARRAY:
470514
length = RANDOM.nextInt(20);
471515
String[] stringArray = new String[length];
@@ -476,7 +520,7 @@ private void fillDataTableWithRandomData(DataTableBuilder dataTableBuilder,
476520
dataTableBuilder.setColumn(colId, stringArray);
477521
break;
478522
default:
479-
break;
523+
throw new UnsupportedOperationException("Unable to generate random data for: " + columnDataTypes[colId]);
480524
}
481525
}
482526
dataTableBuilder.finishRow();
@@ -499,9 +543,21 @@ private void verifyDataIsSame(DataTable newDataTable, DataSchema.ColumnDataType[
499543
case DOUBLE:
500544
Assert.assertEquals(newDataTable.getDouble(rowId, colId), DOUBLES[rowId], ERROR_MESSAGE);
501545
break;
546+
case BIG_DECIMAL:
547+
Assert.assertEquals(newDataTable.getBigDecimal(rowId, colId), BIG_DECIMALS[rowId], ERROR_MESSAGE);
548+
break;
549+
case BOOLEAN:
550+
Assert.assertEquals(newDataTable.getInt(rowId, colId), BOOLEANS[rowId], ERROR_MESSAGE);
551+
break;
552+
case TIMESTAMP:
553+
Assert.assertEquals(newDataTable.getLong(rowId, colId), TIMESTAMPS[rowId], ERROR_MESSAGE);
554+
break;
502555
case STRING:
503556
Assert.assertEquals(newDataTable.getString(rowId, colId), STRINGS[rowId], ERROR_MESSAGE);
504557
break;
558+
case JSON:
559+
Assert.assertEquals(newDataTable.getString(rowId, colId), JSONS[rowId], ERROR_MESSAGE);
560+
break;
505561
case BYTES:
506562
Assert.assertEquals(newDataTable.getBytes(rowId, colId).getBytes(), BYTES[rowId], ERROR_MESSAGE);
507563
break;
@@ -512,8 +568,8 @@ private void verifyDataIsSame(DataTable newDataTable, DataSchema.ColumnDataType[
512568
Assert.assertTrue(Arrays.equals(newDataTable.getIntArray(rowId, colId), INT_ARRAYS[rowId]), ERROR_MESSAGE);
513569
break;
514570
case LONG_ARRAY:
515-
Assert
516-
.assertTrue(Arrays.equals(newDataTable.getLongArray(rowId, colId), LONG_ARRAYS[rowId]), ERROR_MESSAGE);
571+
Assert.assertTrue(Arrays.equals(newDataTable.getLongArray(rowId, colId), LONG_ARRAYS[rowId]),
572+
ERROR_MESSAGE);
517573
break;
518574
case FLOAT_ARRAY:
519575
Assert.assertTrue(Arrays.equals(newDataTable.getFloatArray(rowId, colId), FLOAT_ARRAYS[rowId]),
@@ -523,12 +579,23 @@ private void verifyDataIsSame(DataTable newDataTable, DataSchema.ColumnDataType[
523579
Assert.assertTrue(Arrays.equals(newDataTable.getDoubleArray(rowId, colId), DOUBLE_ARRAYS[rowId]),
524580
ERROR_MESSAGE);
525581
break;
582+
case BOOLEAN_ARRAY:
583+
Assert.assertTrue(Arrays.equals(newDataTable.getIntArray(rowId, colId), BOOLEAN_ARRAYS[rowId]),
584+
ERROR_MESSAGE);
585+
break;
586+
case TIMESTAMP_ARRAY:
587+
Assert.assertTrue(Arrays.equals(newDataTable.getLongArray(rowId, colId), TIMESTAMP_ARRAYS[rowId]),
588+
ERROR_MESSAGE);
589+
break;
590+
case BYTES_ARRAY:
591+
// TODO: add once implementation of datatable bytes array support is added
592+
break;
526593
case STRING_ARRAY:
527594
Assert.assertTrue(Arrays.equals(newDataTable.getStringArray(rowId, colId), STRING_ARRAYS[rowId]),
528595
ERROR_MESSAGE);
529596
break;
530597
default:
531-
break;
598+
throw new UnsupportedOperationException("Unable to generate random data for: " + columnDataTypes[colId]);
532599
}
533600
}
534601
}

0 commit comments

Comments
 (0)