paul-rogers commented on a change in pull request #1870: DRILL-7359: Add support for DICT type in RowSet Framework URL: https://github.com/apache/drill/pull/1870#discussion_r347115170
########## File path: exec/java-exec/src/test/java/org/apache/drill/test/rowSet/test/TestRowSet.java ########## @@ -586,6 +595,490 @@ public void testRepeatedMapStructure() { RowSetUtilities.verify(expected, actual); } + @Test + public void testDictStructure() { + final String dictName = "d"; + + final TupleMetadata schema = new SchemaBuilder() + .add("id", MinorType.INT) + .addDict(dictName, MinorType.INT) + .value(MinorType.VARCHAR) // required int + .resumeSchema() + .buildSchema(); + final ExtendableRowSet rowSet = fixture.rowSet(schema); + final RowSetWriter writer = rowSet.writer(); + + // Dict + // Pick out components and lightly test. (Assumes structure + // tested earlier is still valid, so no need to exhaustively + // test again.) + + assertEquals(ObjectType.ARRAY, writer.column(dictName).type()); + assertTrue(writer.column(dictName).schema().isDict()); + + final ScalarWriter idWriter = writer.column(0).scalar(); + final DictWriter dictWriter = writer.column(1).dict(); + + assertEquals(ValueType.INTEGER, dictWriter.keyType()); + assertEquals(ObjectType.SCALAR, dictWriter.valueType()); + + final ScalarWriter keyWriter = dictWriter.keyWriter(); + final ScalarWriter valueWriter = dictWriter.valueWriter().scalar(); + + assertEquals(ValueType.INTEGER, keyWriter.valueType()); + assertEquals(ValueType.STRING, valueWriter.valueType()); + + // Write data + idWriter.setInt(1); + + keyWriter.setInt(11); + valueWriter.setString("a"); + dictWriter.save(); // Advance to next entry position + keyWriter.setInt(12); + valueWriter.setString("b"); + dictWriter.save(); + writer.save(); + + idWriter.setInt(2); + + keyWriter.setInt(21); + valueWriter.setString("c"); + dictWriter.save(); + writer.save(); + + idWriter.setInt(3); + + keyWriter.setInt(31); + valueWriter.setString("d"); + dictWriter.save(); + keyWriter.setInt(32); + valueWriter.setString("e"); + dictWriter.save(); + writer.save(); + + // Finish the row set and get a reader. + + final SingleRowSet actual = writer.done(); + final RowSetReader reader = actual.reader(); + + // Verify reader structure + + assertEquals(ObjectType.ARRAY, reader.column(dictName).type()); + + final DictReader dictReader = reader.dict(1); + assertEquals(ObjectType.ARRAY, dictReader.type()); + + assertEquals(ValueType.INTEGER, dictReader.keyColumnType()); + assertEquals(ObjectType.SCALAR, dictReader.valueColumnType()); + + // Row 1: get value reader with its position set to entry corresponding to a key + + assertTrue(reader.next()); + assertFalse(dictReader.isNull()); // dict itself is not null + + dictReader.getAsString(); + assertEquals("b", dictReader.getValueReader(12).scalar().getString()); + assertEquals("a", dictReader.getValueReader(11).scalar().getString()); + + // compare entire dict + Map<Object, Object> map = map(11, "a", 12, "b"); + assertEquals(map, dictReader.getObject()); + + // Row 2: get Object representation of value directly + + assertTrue(reader.next()); + assertEquals("c", dictReader.get(21)); + assertTrue(dictReader.getValueReader(22).isNull()); // the dict does not contain an entry with the key + + map = map(21, "c"); + assertEquals(map, dictReader.getObject()); + + // Row 3 + + assertTrue(reader.next()); + + assertEquals("d", dictReader.get(31)); + assertEquals("e", dictReader.get(32)); + + map = map(31, "d", 32, "e"); + assertEquals(map, dictReader.getObject()); + + assertFalse(reader.next()); + + // Verify that the dict accessor's value count was set. + + final DictVector dictVector = (DictVector) actual.container().getValueVector(1).getValueVector(); + assertEquals(3, dictVector.getAccessor().getValueCount()); + + final SingleRowSet expected = fixture.rowSetBuilder(schema) + .addRow(1, objArray(11, "a", 12, "b")) + .addRow(2, objArray(21, "c")) + .addRow(3, objArray(31, "d", 32, "e")) + .build(); + RowSetUtilities.verify(expected, actual); + } + + /** + * Utility method to bootstrap a map object easily. + * + * @param entry key-value sequence + * @return map containing key-value pairs from passed sequence + */ + private Map<Object, Object> map(Object... entry) { + if (entry.length % 2 == 1) { + throw new IllegalArgumentException("Array length should be even."); + } + + Map<Object, Object> map = new HashMap<>(); + for (int i = 0; i < entry.length; i += 2) { + map.put(entry[i], entry[i + 1]); + } + return map; + } + + @Test + public void testDictStructureMapValue() { + final String dictName = "d"; + final int bScale = 1; + + final TupleMetadata schema = new SchemaBuilder() + .add("id", MinorType.INT) + .addDict(dictName, MinorType.INT) + .mapValue() + .add("a", MinorType.INT) + .add("b", MinorType.VARDECIMAL, 8, bScale) + .resumeDict() + .resumeSchema() Review comment: Nit: again, maybe indent from `mapValue()` to `resumeSchema()` to show nested structure. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services