Examples with ScalarWriter org.apache.drill.exec.vector.accessor.ScalarWriter used on opensource projects

Example 46 with ScalarWriter

use of org.apache.drill.exec.vector.accessor.ScalarWriter in project drill by apache.

the class TestRowSet method testTopFixedWidthArray.

/**
 * Test an array of ints (as an example fixed-width type)
 * at the top level of a schema.
 */
@Test
public void testTopFixedWidthArray() {
    final TupleMetadata schema = new SchemaBuilder().add("c", MinorType.INT).addArray("a", MinorType.INT).buildSchema();
    final ExtendableRowSet rs1 = fixture.rowSet(schema);
    final RowSetWriter writer = rs1.writer();
    writer.scalar(0).setInt(10);
    final ScalarWriter array = writer.array(1).scalar();
    array.setInt(100);
    array.setInt(110);
    writer.save();
    writer.scalar(0).setInt(20);
    array.setInt(200);
    array.setInt(120);
    array.setInt(220);
    writer.save();
    writer.scalar(0).setInt(30);
    writer.save();
    final SingleRowSet result = writer.done();
    final RowSetReader reader = result.reader();
    final ArrayReader arrayReader = reader.array(1);
    final ScalarReader elementReader = arrayReader.scalar();
    assertTrue(reader.next());
    assertEquals(10, reader.scalar(0).getInt());
    assertEquals(2, arrayReader.size());
    assertTrue(arrayReader.next());
    assertEquals(100, elementReader.getInt());
    assertTrue(arrayReader.next());
    assertEquals(110, elementReader.getInt());
    assertTrue(reader.next());
    assertEquals(20, reader.scalar(0).getInt());
    assertEquals(3, arrayReader.size());
    assertTrue(arrayReader.next());
    assertEquals(200, elementReader.getInt());
    assertTrue(arrayReader.next());
    assertEquals(120, elementReader.getInt());
    assertTrue(arrayReader.next());
    assertEquals(220, elementReader.getInt());
    assertTrue(reader.next());
    assertEquals(30, reader.scalar(0).getInt());
    assertEquals(0, arrayReader.size());
    assertFalse(reader.next());
    final SingleRowSet rs2 = fixture.rowSetBuilder(schema).addRow(10, intArray(100, 110)).addRow(20, intArray(200, 120, 220)).addRow(30, null).build();
    RowSetUtilities.verify(rs1, rs2);
}

Example 47 with ScalarWriter

use of org.apache.drill.exec.vector.accessor.ScalarWriter in project drill by apache.

the class TestRowSet method testScalarArrayStructure.

/**
 * Test a record with a top level array. The focus here is on the
 * scalar array structure.
 *
 * @throws VectorOverflowException should never occur
 */
@Test
public void testScalarArrayStructure() {
    final TupleMetadata schema = new SchemaBuilder().addArray("a", MinorType.INT).buildSchema();
    final ExtendableRowSet rowSet = fixture.rowSet(schema);
    final RowSetWriter writer = rowSet.writer();
    // Repeated Int
    // Verify the invariants of the "full" and "simple" access paths
    assertEquals(ObjectType.ARRAY, writer.column("a").type());
    assertSame(writer.column("a"), writer.column(0));
    assertSame(writer.array("a"), writer.array(0));
    assertSame(writer.column("a").array(), writer.array("a"));
    assertSame(writer.column(0).array(), writer.array(0));
    assertEquals(ObjectType.SCALAR, writer.column("a").array().entry().type());
    assertEquals(ObjectType.SCALAR, writer.column("a").array().entryType());
    assertSame(writer.array(0).entry().scalar(), writer.array(0).scalar());
    assertEquals(ValueType.INTEGER, writer.array(0).scalar().valueType());
    try {
        writer.column(0).scalar();
        fail();
    } catch (final UnsupportedOperationException e) {
    // Expected
    }
    try {
        writer.column(0).tuple();
        fail();
    } catch (final UnsupportedOperationException e) {
    // Expected
    }
    // Write some data
    final ScalarWriter intWriter = writer.array("a").scalar();
    intWriter.setInt(10);
    intWriter.setInt(11);
    writer.save();
    intWriter.setInt(20);
    intWriter.setInt(21);
    intWriter.setInt(22);
    writer.save();
    intWriter.setInt(30);
    writer.save();
    intWriter.setInt(40);
    intWriter.setInt(41);
    writer.save();
    // Finish the row set and get a reader.
    final SingleRowSet actual = writer.done();
    final RowSetReader reader = actual.reader();
    // Verify the invariants of the "full" and "simple" access paths
    assertEquals(ObjectType.ARRAY, writer.column("a").type());
    assertSame(reader.column("a"), reader.column(0));
    assertSame(reader.array("a"), reader.array(0));
    assertSame(reader.column("a").array(), reader.array("a"));
    assertSame(reader.column(0).array(), reader.array(0));
    assertEquals(ObjectType.SCALAR, reader.column("a").array().entryType());
    assertEquals(ValueType.INTEGER, reader.array(0).scalar().valueType());
    // Read and verify the rows
    final ArrayReader arrayReader = reader.array(0);
    final ScalarReader intReader = arrayReader.scalar();
    assertTrue(reader.next());
    assertFalse(arrayReader.isNull());
    assertEquals(2, arrayReader.size());
    assertTrue(arrayReader.next());
    assertEquals(10, intReader.getInt());
    assertTrue(arrayReader.next());
    assertEquals(11, intReader.getInt());
    assertFalse(arrayReader.next());
    assertTrue(reader.next());
    assertFalse(arrayReader.isNull());
    assertEquals(3, arrayReader.size());
    assertTrue(arrayReader.next());
    assertEquals(20, intReader.getInt());
    assertTrue(arrayReader.next());
    assertEquals(21, intReader.getInt());
    assertTrue(arrayReader.next());
    assertEquals(22, intReader.getInt());
    assertFalse(arrayReader.next());
    assertTrue(reader.next());
    assertFalse(arrayReader.isNull());
    assertEquals(1, arrayReader.size());
    assertTrue(arrayReader.next());
    assertEquals(30, intReader.getInt());
    assertFalse(arrayReader.next());
    assertTrue(reader.next());
    assertFalse(arrayReader.isNull());
    assertEquals(2, arrayReader.size());
    assertTrue(arrayReader.next());
    assertEquals(40, intReader.getInt());
    assertTrue(arrayReader.next());
    assertEquals(41, intReader.getInt());
    assertFalse(arrayReader.next());
    assertFalse(reader.next());
    // Test the above again via the writer and reader
    // utility classes.
    final SingleRowSet expected = fixture.rowSetBuilder(schema).addSingleCol(intArray(10, 11)).addSingleCol(intArray(20, 21, 22)).addSingleCol(intArray(30)).addSingleCol(intArray(40, 41)).build();
    RowSetUtilities.verify(expected, actual);
}

Example 48 with ScalarWriter

use of org.apache.drill.exec.vector.accessor.ScalarWriter in project drill by apache.

the class TestRowSet method testDictStructureMapValue.

@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().buildSchema();
    final ExtendableRowSet rowSet = fixture.rowSet(schema);
    final RowSetWriter writer = rowSet.writer();
    // Dict with Map value
    assertEquals(ObjectType.ARRAY, writer.column(dictName).type());
    final ScalarWriter idWriter = writer.scalar(0);
    final DictWriter dictWriter = writer.column(1).dict();
    assertEquals(ValueType.INTEGER, dictWriter.keyType());
    assertEquals(ObjectType.TUPLE, dictWriter.valueType());
    final ScalarWriter keyWriter = dictWriter.keyWriter();
    final TupleWriter valueWriter = dictWriter.valueWriter().tuple();
    assertEquals(ValueType.INTEGER, keyWriter.valueType());
    ScalarWriter aWriter = valueWriter.scalar("a");
    ScalarWriter bWriter = valueWriter.scalar("b");
    assertEquals(ValueType.INTEGER, aWriter.valueType());
    assertEquals(ValueType.DECIMAL, bWriter.valueType());
    // Write data
    idWriter.setInt(1);
    keyWriter.setInt(11);
    aWriter.setInt(10);
    bWriter.setDecimal(BigDecimal.valueOf(1));
    // advance to next entry position
    dictWriter.save();
    keyWriter.setInt(12);
    aWriter.setInt(11);
    bWriter.setDecimal(BigDecimal.valueOf(2));
    dictWriter.save();
    writer.save();
    idWriter.setInt(2);
    keyWriter.setInt(21);
    aWriter.setInt(20);
    bWriter.setDecimal(BigDecimal.valueOf(3));
    dictWriter.save();
    writer.save();
    idWriter.setInt(3);
    keyWriter.setInt(31);
    aWriter.setInt(30);
    bWriter.setDecimal(BigDecimal.valueOf(4));
    dictWriter.save();
    keyWriter.setInt(32);
    aWriter.setInt(31);
    bWriter.setDecimal(BigDecimal.valueOf(5));
    dictWriter.save();
    keyWriter.setInt(33);
    aWriter.setInt(32);
    bWriter.setDecimal(BigDecimal.valueOf(6));
    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.TUPLE, dictReader.valueColumnType());
    final KeyAccessor keyAccessor = dictReader.keyAccessor();
    final TupleReader valueReader = dictReader.valueReader().tuple();
    // Row 1: get value reader with its position set to entry corresponding to a key
    assertTrue(reader.next());
    // dict itself is not null
    assertFalse(dictReader.isNull());
    assertTrue(keyAccessor.find(12));
    assertEquals(11, valueReader.scalar("a").getInt());
    assertEquals(BigDecimal.valueOf(2.0), valueReader.scalar("b").getDecimal());
    // MapReader#getObject() returns a List containing values for each column
    // rather than mapping of column name to it's value, hence List is expected for Dict's value.
    Map<Object, Object> map = map(11, Arrays.asList(10, BigDecimal.valueOf(1.0)), 12, Arrays.asList(11, BigDecimal.valueOf(2.0)));
    assertEquals(map, dictReader.getObject());
    // Row 2
    assertTrue(reader.next());
    assertFalse(keyAccessor.find(222));
    assertTrue(keyAccessor.find(21));
    assertEquals(Arrays.asList(20, BigDecimal.valueOf(3.0)), valueReader.getObject());
    map = map(21, Arrays.asList(20, BigDecimal.valueOf(3.0)));
    assertEquals(map, dictReader.getObject());
    // Row 3
    assertTrue(reader.next());
    assertTrue(keyAccessor.find(32));
    assertFalse(valueReader.isNull());
    assertEquals(31, valueReader.scalar("a").getInt());
    assertEquals(BigDecimal.valueOf(5.0), valueReader.scalar("b").getDecimal());
    assertTrue(keyAccessor.find(31));
    assertEquals(30, valueReader.scalar("a").getInt());
    assertEquals(BigDecimal.valueOf(4.0), valueReader.scalar("b").getDecimal());
    assertFalse(keyAccessor.find(404));
    map = map(31, Arrays.asList(30, BigDecimal.valueOf(4.0)), 32, Arrays.asList(31, BigDecimal.valueOf(5.0)), 33, Arrays.asList(32, BigDecimal.valueOf(6.0)));
    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, map(11, objArray(10, BigDecimal.valueOf(1.0)), 12, objArray(11, BigDecimal.valueOf(2.0)))).addRow(2, map(21, objArray(20, BigDecimal.valueOf(3.0)))).addRow(3, map(31, objArray(30, BigDecimal.valueOf(4.0)), 32, objArray(31, BigDecimal.valueOf(5.0)), 33, objArray(32, BigDecimal.valueOf(6.0)))).build();
    RowSetUtilities.verify(expected, actual);
}

Example 49 with ScalarWriter

use of org.apache.drill.exec.vector.accessor.ScalarWriter in project drill by apache.

the class TestResultSetLoaderMaps method testMapOverflowWithNewColumn.

/**
 * Test the case in which a new column is added during the overflow row. Unlike
 * the top-level schema case, internally we must create a copy of the map, and
 * move vectors across only when the result is to include the schema version
 * of the target column. For overflow, the new column is added after the
 * first batch; it is added in the second batch that contains the overflow
 * row in which the column was added.
 */
@Test
public void testMapOverflowWithNewColumn() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").add("b", MinorType.INT).add("c", MinorType.VARCHAR).resumeSchema().buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).rowCountLimit(ValueVector.MAX_ROW_COUNT).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    assertEquals(4, rsLoader.schemaVersion());
    final RowSetLoader rootWriter = rsLoader.writer();
    // Can't use the shortcut to populate rows when doing a schema
    // change.
    final ScalarWriter aWriter = rootWriter.scalar("a");
    final TupleWriter mWriter = rootWriter.tuple("m");
    final ScalarWriter bWriter = mWriter.scalar("b");
    final ScalarWriter cWriter = mWriter.scalar("c");
    final byte[] value = new byte[512];
    Arrays.fill(value, (byte) 'X');
    int count = 0;
    rsLoader.startBatch();
    while (!rootWriter.isFull()) {
        rootWriter.start();
        aWriter.setInt(count);
        bWriter.setInt(count * 10);
        cWriter.setBytes(value, value.length);
        if (rootWriter.isFull()) {
            // Overflow just occurred. Add another column.
            mWriter.addColumn(SchemaBuilder.columnSchema("d", MinorType.INT, DataMode.OPTIONAL));
            mWriter.scalar("d").setInt(count * 100);
        }
        rootWriter.save();
        count++;
    }
    // Result set should include the original columns, but not d.
    RowSet result = fixture.wrap(rsLoader.harvest());
    assertEquals(4, rsLoader.schemaVersion());
    assertTrue(schema.isEquivalent(result.schema()));
    final BatchSchema expectedSchema = new BatchSchema(SelectionVectorMode.NONE, schema.toFieldList());
    assertTrue(expectedSchema.isEquivalent(result.batchSchema()));
    // Use a reader to validate row-by-row. Too large to create an expected
    // result set.
    RowSetReader reader = result.reader();
    TupleReader mapReader = reader.tuple("m");
    int rowId = 0;
    while (reader.next()) {
        assertEquals(rowId, reader.scalar("a").getInt());
        assertEquals(rowId * 10, mapReader.scalar("b").getInt());
        assertTrue(Arrays.equals(value, mapReader.scalar("c").getBytes()));
        rowId++;
    }
    result.clear();
    // Next batch should start with the overflow row
    rsLoader.startBatch();
    assertEquals(1, rootWriter.rowCount());
    result = fixture.wrap(rsLoader.harvest());
    assertEquals(1, result.rowCount());
    reader = result.reader();
    mapReader = reader.tuple("m");
    while (reader.next()) {
        assertEquals(rowId, reader.scalar("a").getInt());
        assertEquals(rowId * 10, mapReader.scalar("b").getInt());
        assertTrue(Arrays.equals(value, mapReader.scalar("c").getBytes()));
        assertEquals(rowId * 100, mapReader.scalar("d").getInt());
    }
    result.clear();
    rsLoader.close();
}

Example 50 with ScalarWriter

use of org.apache.drill.exec.vector.accessor.ScalarWriter in project drill by apache.

the class TestResultSetLoaderOverflow method testMissingArrayValues.

/**
 * Test the case that an array has "missing values" before the overflow.
 */
@Test
public void testMissingArrayValues() {
    TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).add("b", MinorType.VARCHAR).addArray("c", MinorType.INT).buildSchema();
    ResultSetOptions options = new ResultSetOptionBuilder().rowCountLimit(ValueVector.MAX_ROW_COUNT).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    byte[] value = new byte[512];
    Arrays.fill(value, (byte) 'X');
    int blankAfter = ValueVector.MAX_BUFFER_SIZE / 512 * 2 / 3;
    ScalarWriter cWriter = rootWriter.array("c").scalar();
    rsLoader.startBatch();
    int rowId = 0;
    while (rootWriter.start()) {
        rootWriter.scalar("a").setInt(rowId);
        rootWriter.scalar("b").setBytes(value, value.length);
        if (rowId < blankAfter) {
            for (int i = 0; i < 3; i++) {
                cWriter.setInt(rowId * 3 + i);
            }
        }
        rootWriter.save();
        rowId++;
    }
    VectorContainer container = rsLoader.harvest();
    BatchValidator.validate(container);
    RowSet result = fixture.wrap(container);
    assertEquals(rowId - 1, result.rowCount());
    RowSetReader reader = result.reader();
    ArrayReader cArray = reader.array("c");
    ScalarReader cReader = cArray.scalar();
    while (reader.next()) {
        assertEquals(reader.offset(), reader.scalar("a").getInt());
        assertArrayEquals(value, reader.scalar("b").getBytes());
        if (reader.offset() < blankAfter) {
            assertEquals(3, cArray.size());
            for (int i = 0; i < 3; i++) {
                assertTrue(cArray.next());
                assertEquals(reader.offset() * 3 + i, cReader.getInt());
            }
        } else {
            assertEquals(0, cArray.size());
        }
    }
    result.clear();
    rsLoader.close();
}