Examples with BsonDocumentReader org.bson.BsonDocumentReader used on opensource projects

Example 16 with BsonDocumentReader

use of org.bson.BsonDocumentReader in project drill by apache.

the class TestBsonRecordReader method testObjectIdType.

@Test
public void testObjectIdType() throws IOException {
    BsonDocument bsonDoc = new BsonDocument();
    BsonObjectId value = new BsonObjectId(new ObjectId());
    bsonDoc.append("_idKey", value);
    writer.reset();
    bsonReader.write(writer, new BsonDocumentReader(bsonDoc));
    SingleMapReaderImpl mapReader = (SingleMapReaderImpl) writer.getMapVector().getReader();
    byte[] readByteArray = mapReader.reader("_idKey").readByteArray();
    assertArrayEquals(value.getValue().toByteArray(), readByteArray);
}

Example 17 with BsonDocumentReader

use of org.bson.BsonDocumentReader in project drill by apache.

the class TestBsonRecordReader method testArrayOfDocumentType.

@Test
public void testArrayOfDocumentType() throws IOException {
    BsonDocument bsonDoc = new BsonDocument();
    BsonWriter bw = new BsonDocumentWriter(bsonDoc);
    bw.writeStartDocument();
    bw.writeName("a");
    bw.writeString("MongoDB");
    bw.writeName("b");
    bw.writeStartArray();
    bw.writeStartDocument();
    bw.writeName("c");
    bw.writeInt32(1);
    bw.writeEndDocument();
    bw.writeEndArray();
    bw.writeEndDocument();
    bw.flush();
    writer.reset();
    bsonReader.write(writer, new BsonDocumentReader(bsonDoc));
    FieldReader reader = writer.getMapVector().getReader();
    SingleMapReaderImpl mapReader = (SingleMapReaderImpl) reader;
    FieldReader reader3 = mapReader.reader("b");
    assertEquals("MongoDB", mapReader.reader("a").readText().toString());
}

Example 18 with BsonDocumentReader

use of org.bson.BsonDocumentReader in project drill by apache.

the class TestBsonRecordReader method testBooleanType.

@Test
public void testBooleanType() throws IOException {
    BsonDocument bsonDoc = new BsonDocument();
    bsonDoc.append("booleanKey", new BsonBoolean(true));
    writer.reset();
    bsonReader.write(writer, new BsonDocumentReader(bsonDoc));
    SingleMapReaderImpl mapReader = (SingleMapReaderImpl) writer.getMapVector().getReader();
    assertTrue(mapReader.reader("booleanKey").readBoolean());
}

Example 19 with BsonDocumentReader

use of org.bson.BsonDocumentReader in project drill by apache.

the class TestBsonRecordReader method testTimeStampType.

@Test
public void testTimeStampType() throws IOException {
    BsonDocument bsonDoc = new BsonDocument();
    bsonDoc.append("ts_small", new BsonTimestamp(1000, 10));
    bsonDoc.append("ts_large", new BsonTimestamp(1000000000, 10));
    writer.reset();
    bsonReader.write(writer, new BsonDocumentReader(bsonDoc));
    SingleMapReaderImpl mapReader = (SingleMapReaderImpl) writer.getMapVector().getReader();
    assertEquals(1000000L, mapReader.reader("ts_small").readLocalDateTime().atZone(ZoneOffset.systemDefault()).toInstant().toEpochMilli());
    assertEquals(1000000000000L, mapReader.reader("ts_large").readLocalDateTime().atZone(ZoneOffset.systemDefault()).toInstant().toEpochMilli());
}

Example 20 with BsonDocumentReader

use of org.bson.BsonDocumentReader in project drill by apache.

the class MongoRecordReader method next.

@Override
public int next() {
    if (cursor == null) {
        logger.debug("Filters Applied : " + filters);
        logger.debug("Fields Selected :" + fields);
        MongoIterable<BsonDocument> projection;
        if (CollectionUtils.isNotEmpty(operations)) {
            List<Bson> operations = new ArrayList<>(this.operations);
            if (!fields.isEmpty()) {
                operations.add(Aggregates.project(fields));
            }
            if (plugin.getConfig().allowDiskUse()) {
                projection = collection.aggregate(operations).allowDiskUse(true);
            } else {
                projection = collection.aggregate(operations);
            }
        } else {
            projection = collection.find(filters).projection(fields);
        }
        cursor = projection.batchSize(plugin.getConfig().getBatchSize()).iterator();
    }
    writer.allocate();
    writer.reset();
    int docCount = 0;
    Stopwatch watch = Stopwatch.createStarted();
    try {
        while (docCount < BaseValueVector.INITIAL_VALUE_ALLOCATION && cursor.hasNext()) {
            writer.setPosition(docCount);
            if (isBsonRecordReader) {
                BsonDocument bsonDocument = cursor.next();
                bsonReader.write(writer, new BsonDocumentReader(bsonDocument));
            } else {
                String doc = cursor.next().toJson();
                jsonReader.setSource(doc.getBytes(Charsets.UTF_8));
                jsonReader.write(writer);
            }
            docCount++;
        }
        if (isBsonRecordReader) {
            bsonReader.ensureAtLeastOneField(writer);
        } else {
            jsonReader.ensureAtLeastOneField(writer);
        }
        writer.setValueCount(docCount);
        logger.debug("Took {} ms to get {} records", watch.elapsed(TimeUnit.MILLISECONDS), docCount);
        return docCount;
    } catch (IOException e) {
        String msg = "Failure while reading document. - Parser was at record: " + (docCount + 1);
        logger.error(msg, e);
        throw new DrillRuntimeException(msg, e);
    }
}