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

Example 31 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader 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 32 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader 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 33 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader 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 34 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader 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();
}

Example 35 with RowSetLoader

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

the class BaseMockBatchReader method makeBatch.

protected void makeBatch() {
    RowSetLoader writer = tableLoader.writer();
    int offset = (batchCount - 1) * 20 + startIndex;
    writeRow(writer, offset + 10, "fred");
    writeRow(writer, offset + 20, "wilma");
}