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

Example 51 with RowSetLoader

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

the class TestResultSetLoaderProjection method testMapProjectionMemberAndMap.

@Test
public void testMapProjectionMemberAndMap() {
    // SELECT m1, m1.b
    // This really means project all of m1; m1.b is along for the ride.
    List<SchemaPath> selection = RowSetTestUtils.projectList("m1", "m1.b");
    // Define an "early" reader schema consistent with the projection.
    TupleMetadata schema = new SchemaBuilder().addMap("m1").add("a", MinorType.INT).add("b", 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());
    // Write a couple of rows.
    rsLoader.startBatch();
    rootWriter.start();
    rootWriter.addSingleCol(mapValue(1, 2)).addSingleCol(mapValue(11, 12));
    // Verify. The whole map appears in the result set because the
    // project list included the whole map as well as a map member.
    SingleRowSet expected = fixture.rowSetBuilder(schema).addSingleCol(mapValue(1, 2)).addSingleCol(mapValue(11, 12)).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
    rsLoader.close();
}

Example 52 with RowSetLoader

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

the class TestResultSetLoaderProjection method setupProvidedSchema.

// Setup to test the various project/provided schema cases in
// ProjectionFilter, especially CompoundProjectionFilter
public ResultSetLoader setupProvidedSchema(boolean isStrict, List<SchemaPath> selection) {
    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();
    // Provided schema: disjoint set of above.
    TupleMetadata providedSchema = new SchemaBuilder().addMap(// Same
    "m1").add("a", MinorType.INT).add("b", MinorType.INT).add("z", // Add a column
    MinorType.INT).resumeSchema().addMap(// Omit c
    "m2").add("d", MinorType.INT).resumeSchema().addMap(// Add m4
    "m4").add("g", MinorType.INT).add("h", MinorType.INT).resumeSchema().build();
    if (isStrict) {
        SchemaUtils.markStrict(providedSchema);
    }
    RequestedTuple proj = Projections.parse(selection);
    CustomErrorContext errorContext = new EmptyErrorContext();
    ProjectionFilter projectionFilter = ProjectionFilter.providedSchemaFilter(proj, providedSchema, errorContext);
    ResultSetOptions options = new ResultSetOptionBuilder().projectionFilter(projectionFilter).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    // Write a couple of rows.
    rsLoader.startBatch();
    RowSetLoader rootWriter = rsLoader.writer();
    rootWriter.start();
    rootWriter.addRow(mapValue(1, 2), mapValue(3, 4), mapValue(5, 6)).addRow(mapValue(11, 12), mapValue(13, 14), mapValue(15, 16));
    return rsLoader;
}

Example 53 with RowSetLoader

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

the class TestResultSetLoaderRepeatedList method test2DLateSchemaIncremental.

@Test
public void test2DLateSchemaIncremental() {
    final TupleMetadata schema = new SchemaBuilder().add("id", MinorType.INT).addRepeatedList("list1").addArray(MinorType.VARCHAR).resumeSchema().addRepeatedList("list2").addArray(MinorType.VARCHAR).resumeSchema().buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    final RowSetLoader writer = rsLoader.writer();
    // Add columns dynamically
    writer.addColumn(schema.metadata(0));
    // Write a row without the array.
    rsLoader.startBatch();
    writer.addRow(1);
    // Add the repeated list, but without contents.
    writer.addColumn(schema.metadata(1).cloneEmpty());
    // Sanity check of writer structure
    assertEquals(2, writer.size());
    final ObjectWriter listObj = writer.column("list1");
    assertEquals(ObjectType.ARRAY, listObj.type());
    final ArrayWriter listWriter = listObj.array();
    // No child defined yet. A dummy child is inserted instead.
    assertEquals(MinorType.NULL, listWriter.entry().schema().type());
    assertEquals(ObjectType.ARRAY, listWriter.entryType());
    assertEquals(ObjectType.SCALAR, listWriter.array().entryType());
    assertEquals(ValueType.NULL, listWriter.array().scalar().valueType());
    // Although we don't know the type of the inner, we can still
    // create null (empty) elements in the outer array.
    writer.addRow(2, null).addRow(3, objArray()).addRow(4, objArray(objArray(), null));
    // Define the inner type.
    final RepeatedListWriter listWriterImpl = (RepeatedListWriter) listWriter;
    listWriterImpl.defineElement(MaterializedField.create("list1", Types.repeated(MinorType.VARCHAR)));
    // Sanity check of completed structure
    assertEquals(ObjectType.ARRAY, listWriter.entryType());
    final ArrayWriter innerWriter = listWriter.array();
    assertEquals(ObjectType.SCALAR, innerWriter.entryType());
    final ScalarWriter strWriter = innerWriter.scalar();
    assertEquals(ValueType.STRING, strWriter.valueType());
    // Write values
    writer.addRow(5, objArray(strArray("a1", "b1"), strArray("c1", "d1")));
    // Add the second list, with a complete type
    writer.addColumn(schema.metadata(2));
    // Sanity check of writer structure
    assertEquals(3, writer.size());
    final ObjectWriter list2Obj = writer.column("list2");
    assertEquals(ObjectType.ARRAY, list2Obj.type());
    final ArrayWriter list2Writer = list2Obj.array();
    assertEquals(ObjectType.ARRAY, list2Writer.entryType());
    final ArrayWriter inner2Writer = list2Writer.array();
    assertEquals(ObjectType.SCALAR, inner2Writer.entryType());
    final ScalarWriter str2Writer = inner2Writer.scalar();
    assertEquals(ValueType.STRING, str2Writer.valueType());
    // Write values
    writer.addRow(6, objArray(strArray("a2", "b2"), strArray("c2", "d2")), objArray(strArray("w2", "x2"), strArray("y2", "z2")));
    // Add the second list, with a complete type
    // Verify the values.
    // (Relies on the row set level repeated list tests having passed.)
    final RowSet expected = fixture.rowSetBuilder(schema).addRow(1, objArray(), objArray()).addRow(2, objArray(), objArray()).addRow(3, objArray(), objArray()).addRow(4, objArray(objArray(), null), objArray()).addRow(5, objArray(strArray("a1", "b1"), strArray("c1", "d1")), objArray()).addRow(6, objArray(strArray("a2", "b2"), strArray("c2", "d2")), objArray(strArray("w2", "x2"), strArray("y2", "z2"))).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
    rsLoader.close();
}

Example 54 with RowSetLoader

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

the class TestResultSetLoaderDictArray method testKeyOverflow.

@Test
public void testKeyOverflow() {
    TupleMetadata schema = new SchemaBuilder().addDictArray("d", MinorType.VARCHAR).value(MinorType.INT).resumeSchema().buildSchema();
    ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().rowCountLimit(ValueVector.MAX_ROW_COUNT).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    rsLoader.startBatch();
    byte[] key = new byte[523];
    Arrays.fill(key, (byte) 'X');
    // number of dicts in each row
    int arraySize = 3;
    // number of entries in each dict
    int dictSize = 1;
    // Number of rows should be driven by vector size.
    // Our row count should include the overflow row
    ArrayWriter arrayDictWriter = rootWriter.array(0);
    DictWriter dictWriter = arrayDictWriter.dict();
    ScalarWriter keyWriter = dictWriter.keyWriter();
    ScalarWriter valueWriter = dictWriter.valueWriter().scalar();
    int expectedCount = ValueVector.MAX_BUFFER_SIZE / (key.length * dictSize * arraySize);
    // System.out.println("expectedCoutn: " + expectedCount);
    {
        int count = 0;
        while (!rootWriter.isFull()) {
            rootWriter.start();
            for (int i = 0; i < arraySize; i++) {
                for (int j = 0; j < dictSize; j++) {
                    keyWriter.setBytes(key, key.length);
                    // acts as a placeholder, the actual value is not important
                    valueWriter.setInt(0);
                    // not necessary for scalars, just for completeness
                    dictWriter.save();
                }
                arrayDictWriter.save();
            }
            rootWriter.save();
            count++;
        }
        assertEquals(expectedCount + 1, count);
        // System.out.println("count: " + count);
        // Loader's row count should include only "visible" rows
        assertEquals(expectedCount, rootWriter.rowCount());
        // Total count should include invisible and look-ahead rows.
        assertEquals(expectedCount + 1, rsLoader.totalRowCount());
        // Result should exclude the overflow row
        VectorContainer container = rsLoader.harvest();
        BatchValidator.validate(container);
        RowSet result = fixture.wrap(container);
        assertEquals(expectedCount, result.rowCount());
        result.clear();
    }
    // Next batch should start with the overflow row
    {
        rsLoader.startBatch();
        assertEquals(1, rootWriter.rowCount());
        assertEquals(expectedCount + 1, rsLoader.totalRowCount());
        VectorContainer container = rsLoader.harvest();
        BatchValidator.validate(container);
        RowSet result = fixture.wrap(container);
        assertEquals(1, result.rowCount());
        result.clear();
    }
    rsLoader.close();
}

Example 55 with RowSetLoader

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

the class TestResultSetLoaderDictArray method testBasics.

@Test
public void testBasics() {
    TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addDictArray("d", MinorType.INT).value(MinorType.VARCHAR).resumeSchema().buildSchema();
    ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    // Verify structure and schema
    TupleMetadata actualSchema = rootWriter.tupleSchema();
    assertEquals(2, actualSchema.size());
    assertTrue(actualSchema.metadata(1).isArray());
    assertTrue(actualSchema.metadata(1).isDict());
    assertEquals(2, actualSchema.metadata("d").tupleSchema().size());
    assertEquals(2, actualSchema.column("d").getChildren().size());
    DictWriter dictWriter = rootWriter.array("d").dict();
    assertSame(actualSchema.metadata("d").tupleSchema(), dictWriter.schema().tupleSchema());
    // Write a couple of rows with arrays.
    rsLoader.startBatch();
    rootWriter.addRow(10, objArray(map(110, "d1.1", 111, "d1.2", 112, "d1.3"), map(120, "d2.2"))).addRow(20, objArray()).addRow(30, objArray(map(310, "d3.1", 311, "d3.2", 313, "d3.4", 317, "d3.9"), map(320, "d4.2"), map(332, "d5.1", 339, "d5.5", 337, "d5.6")));
    // Verify the batch
    RowSet actual = fixture.wrap(rsLoader.harvest());
    RepeatedDictVector repeatedDictVector = (RepeatedDictVector) actual.container().getValueVector(1).getValueVector();
    // RepeatedDictVector contains one child - DictVector
    MaterializedField dictArrayField = repeatedDictVector.getField();
    assertEquals(1, dictArrayField.getChildren().size());
    DictVector dictVector = (DictVector) repeatedDictVector.getDataVector();
    Iterator<MaterializedField> iter = dictVector.getField().getChildren().iterator();
    assertTrue(dictWriter.keyWriter().schema().schema().isEquivalent(iter.next()));
    assertTrue(dictWriter.valueWriter().scalar().schema().schema().isEquivalent(iter.next()));
    SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(10, objArray(map(110, "d1.1", 111, "d1.2", 112, "d1.3"), map(120, "d2.2"))).addRow(20, objArray()).addRow(30, objArray(map(310, "d3.1", 311, "d3.2", 313, "d3.4", 317, "d3.9"), map(320, "d4.2"), map(332, "d5.1", 339, "d5.5", 337, "d5.6"))).build();
    RowSetUtilities.verify(expected, actual);
    rsLoader.close();
}