Examples with ResultSetLoader org.apache.drill.exec.physical.resultSet.ResultSetLoader used on opensource projects

Example 46 with ResultSetLoader

use of org.apache.drill.exec.physical.resultSet.ResultSetLoader in project drill by apache.

the class TestResultSetLoaderMaps method testMapAddition.

/**
 * Test adding a map to a loader after writing the first row.
 */
@Test
public void testMapAddition() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    assertEquals(1, rsLoader.schemaVersion());
    final RowSetLoader rootWriter = rsLoader.writer();
    // Start without the map. Add a map after the first row.
    rsLoader.startBatch();
    rootWriter.addRow(10);
    final int mapIndex = rootWriter.addColumn(SchemaBuilder.columnSchema("m", MinorType.MAP, DataMode.REQUIRED));
    final TupleWriter mapWriter = rootWriter.tuple(mapIndex);
    // Add a column to the map with the same name as the top-level column.
    // Verifies that the name spaces are independent.
    final int colIndex = mapWriter.addColumn(SchemaBuilder.columnSchema("a", MinorType.VARCHAR, DataMode.REQUIRED));
    assertEquals(0, colIndex);
    // Ensure metadata was added
    assertTrue(mapWriter.tupleSchema().size() == 1);
    assertSame(mapWriter.tupleSchema(), mapWriter.schema().tupleSchema());
    assertSame(mapWriter.tupleSchema().metadata(colIndex), mapWriter.scalar(colIndex).schema());
    rootWriter.addRow(20, mapValue("fred")).addRow(30, mapValue("barney"));
    final RowSet actual = fixture.wrap(rsLoader.harvest());
    assertEquals(3, rsLoader.schemaVersion());
    assertEquals(3, actual.rowCount());
    final MapVector mapVector = (MapVector) actual.container().getValueVector(1).getValueVector();
    final MaterializedField mapField = mapVector.getField();
    assertEquals(1, mapField.getChildren().size());
    assertTrue(mapWriter.scalar(colIndex).schema().schema().isEquivalent(mapField.getChildren().iterator().next()));
    // Validate first batch
    final TupleMetadata expectedSchema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").add("a", MinorType.VARCHAR).resumeSchema().buildSchema();
    final SingleRowSet expected = fixture.rowSetBuilder(expectedSchema).addRow(10, mapValue("")).addRow(20, mapValue("fred")).addRow(30, mapValue("barney")).build();
    RowSetUtilities.verify(expected, actual);
    rsLoader.close();
}

Example 47 with ResultSetLoader

use of org.apache.drill.exec.physical.resultSet.ResultSetLoader in project drill by apache.

the class TestResultSetLoaderProjection method testArrayProjection.

@Test
public void testArrayProjection() {
    List<SchemaPath> selection = RowSetTestUtils.projectList("a1", "a2[0]");
    TupleMetadata schema = new SchemaBuilder().addArray("a1", MinorType.INT).addArray("a2", MinorType.INT).addArray("a3", MinorType.INT).buildSchema();
    ResultSetOptions options = new ResultSetOptionBuilder().projection(Projections.parse(selection)).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    // Verify the projected columns
    TupleMetadata actualSchema = rootWriter.tupleSchema();
    assertTrue(actualSchema.metadata("a1").isArray());
    assertTrue(rootWriter.column("a1").isProjected());
    assertTrue(actualSchema.metadata("a2").isArray());
    assertTrue(rootWriter.column("a2").isProjected());
    assertTrue(actualSchema.metadata("a3").isArray());
    assertFalse(rootWriter.column("a3").isProjected());
    // Write a couple of rows.
    rsLoader.startBatch();
    rootWriter.start();
    rootWriter.addRow(intArray(10, 100), intArray(20, 200), intArray(30, 300)).addRow(intArray(11, 101), intArray(21, 201), intArray(31, 301));
    // Verify. Only the projected columns appear in the result set.
    TupleMetadata expectedSchema = new SchemaBuilder().addArray("a1", MinorType.INT).addArray("a2", MinorType.INT).buildSchema();
    SingleRowSet expected = fixture.rowSetBuilder(expectedSchema).addRow(intArray(10, 100), intArray(20, 200)).addRow(intArray(11, 101), intArray(21, 201)).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
    rsLoader.close();
}

Example 48 with ResultSetLoader

use of org.apache.drill.exec.physical.resultSet.ResultSetLoader in project drill by apache.

the class TestResultSetLoaderProjection method testStrictFullMapProjectionWithSchema.

/**
 * Projection is based on both the projection list and the
 * provided schema, if strict.
 */
@Test
public void testStrictFullMapProjectionWithSchema() {
    // From the reader's perspective, m1, m2 and m3 are projected, m4 is not
    // (the reader does not ask m4 to be created)
    // m1.z is not requested, m2.c is, but is not projected.
    // None of m3 is projected.
    List<SchemaPath> selection = RowSetTestUtils.projectList("m1", "m2", "m3", "m4");
    ResultSetLoader rsLoader = setupProvidedSchema(true, selection);
    RowSetLoader rootWriter = rsLoader.writer();
    // Verify the projected columns
    TupleMetadata actualSchema = rootWriter.tupleSchema();
    TupleWriter m1Writer = rootWriter.tuple("m1");
    assertTrue(m1Writer.isProjected());
    assertEquals(2, m1Writer.tupleSchema().size());
    assertTrue(m1Writer.column("a").isProjected());
    assertTrue(m1Writer.column("b").isProjected());
    TupleWriter m2Writer = rootWriter.tuple("m2");
    assertTrue(m2Writer.isProjected());
    assertEquals(2, m2Writer.tupleSchema().size());
    assertFalse(m2Writer.column("c").isProjected());
    assertTrue(m2Writer.column("d").isProjected());
    TupleWriter m3Writer = rootWriter.tuple("m3");
    assertFalse(m3Writer.isProjected());
    assertEquals(2, m3Writer.tupleSchema().size());
    assertFalse(m3Writer.column("e").isProjected());
    assertFalse(m3Writer.column("f").isProjected());
    assertNull(actualSchema.metadata("m4"));
    // Verify. Only the projected columns appear in the result set.
    TupleMetadata expectedSchema = new SchemaBuilder().addMap("m1").add("a", MinorType.INT).add("b", MinorType.INT).resumeSchema().addMap("m2").add("d", MinorType.INT).resumeSchema().build();
    SingleRowSet expected = fixture.rowSetBuilder(expectedSchema).addRow(mapValue(1, 2), mapValue(4)).addRow(mapValue(11, 12), mapValue(14)).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
    rsLoader.close();
}

Example 49 with ResultSetLoader

use of org.apache.drill.exec.physical.resultSet.ResultSetLoader in project drill by apache.

the class TestResultSetLoaderProjection method testMapProjection.

@Test
public void testMapProjection() {
    List<SchemaPath> selection = RowSetTestUtils.projectList("m1", "m2.d");
    TupleMetadata schema = new SchemaBuilder().addMap("m1").add("a", MinorType.INT).add("b", MinorType.INT).resumeSchema().addMap("m2").add("c", MinorType.INT).add("d", MinorType.INT).resumeSchema().addMap("m3").add("e", MinorType.INT).add("f", MinorType.INT).resumeSchema().buildSchema();
    ResultSetOptions options = new ResultSetOptionBuilder().projection(Projections.parse(selection)).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    // Verify the projected columns
    TupleMetadata actualSchema = rootWriter.tupleSchema();
    ColumnMetadata m1Md = actualSchema.metadata("m1");
    TupleWriter m1Writer = rootWriter.tuple("m1");
    assertTrue(m1Md.isMap());
    assertTrue(m1Writer.isProjected());
    assertEquals(2, m1Md.tupleSchema().size());
    assertTrue(m1Writer.column("a").isProjected());
    assertTrue(m1Writer.column("b").isProjected());
    ColumnMetadata m2Md = actualSchema.metadata("m2");
    TupleWriter m2Writer = rootWriter.tuple("m2");
    assertTrue(m2Md.isMap());
    assertTrue(m2Writer.isProjected());
    assertEquals(2, m2Md.tupleSchema().size());
    assertFalse(m2Writer.column("c").isProjected());
    assertTrue(m2Writer.column("d").isProjected());
    ColumnMetadata m3Md = actualSchema.metadata("m3");
    TupleWriter m3Writer = rootWriter.tuple("m3");
    assertTrue(m3Md.isMap());
    assertFalse(m3Writer.isProjected());
    assertEquals(2, m3Md.tupleSchema().size());
    assertFalse(m3Writer.column("e").isProjected());
    assertFalse(m3Writer.column("f").isProjected());
    // Write a couple of rows.
    rsLoader.startBatch();
    rootWriter.start();
    rootWriter.addRow(mapValue(1, 2), mapValue(3, 4), mapValue(5, 6)).addRow(mapValue(11, 12), mapValue(13, 14), mapValue(15, 16));
    // Verify. Only the projected columns appear in the result set.
    TupleMetadata expectedSchema = new SchemaBuilder().addMap("m1").add("a", MinorType.INT).add("b", MinorType.INT).resumeSchema().addMap("m2").add("d", MinorType.INT).resumeSchema().buildSchema();
    SingleRowSet expected = fixture.rowSetBuilder(expectedSchema).addRow(mapValue(1, 2), mapValue(4)).addRow(mapValue(11, 12), mapValue(14)).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
    rsLoader.close();
}

Example 50 with ResultSetLoader

use of org.apache.drill.exec.physical.resultSet.ResultSetLoader in project drill by apache.

the class TestResultSetLoaderProjection method testDictProjection.

@Test
public void testDictProjection() {
    final String dictName1 = "d1";
    final String dictName2 = "d2";
    // There is no test for case when obtaining a value by key as this is not as simple projection
    // as it is in case of map - there is a need to find a value corresponding to a key
    // (the functionality is currently present in DictReader) and final column schema should be
    // changed from dict structure with `key` and `value` children to a simple `value`.
    List<SchemaPath> selection = RowSetTestUtils.projectList(dictName1);
    TupleMetadata schema = new SchemaBuilder().addDict(dictName1, MinorType.VARCHAR).value(MinorType.INT).resumeSchema().addDict(dictName2, MinorType.VARCHAR).value(MinorType.INT).resumeSchema().buildSchema();
    ResultSetOptions options = new ResultSetOptionBuilder().projection(Projections.parse(selection)).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    // Verify the projected columns
    TupleMetadata actualSchema = rootWriter.tupleSchema();
    ColumnMetadata dictMetadata1 = actualSchema.metadata(dictName1);
    DictWriter dictWriter1 = rootWriter.dict(dictName1);
    assertTrue(dictMetadata1.isDict());
    assertTrue(dictWriter1.isProjected());
    assertEquals(2, dictMetadata1.tupleSchema().size());
    assertTrue(dictWriter1.keyWriter().isProjected());
    assertTrue(dictWriter1.valueWriter().isProjected());
    ColumnMetadata dictMetadata2 = actualSchema.metadata(dictName2);
    DictWriter dictWriter2 = rootWriter.dict(dictName2);
    assertTrue(dictMetadata2.isDict());
    assertFalse(dictWriter2.isProjected());
    assertEquals(2, dictMetadata2.tupleSchema().size());
    assertFalse(dictWriter2.keyWriter().isProjected());
    assertFalse(dictWriter2.valueWriter().isProjected());
    // Write a couple of rows.
    rsLoader.startBatch();
    rootWriter.start();
    rootWriter.addRow(map("a", 1, "b", 2), map("c", 3, "d", 4)).addRow(map("a", 11, "b", 12), map("c", 13, "d", 14));
    // Verify. Only the projected columns appear in the result set.
    TupleMetadata expectedSchema = new SchemaBuilder().addDict(dictName1, MinorType.VARCHAR).value(MinorType.INT).resumeSchema().buildSchema();
    SingleRowSet expected = fixture.rowSetBuilder(expectedSchema).addRow(map("a", 1, "b", 2)).addRow(map("a", 11, "b", 12)).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
    rsLoader.close();
}