Java tutorial
/* * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.facebook.presto.mongodb; import com.facebook.presto.spi.ConnectorPageSink; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.Page; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.spi.StandardErrorCode; import com.facebook.presto.spi.block.Block; import com.facebook.presto.spi.type.BigintType; import com.facebook.presto.spi.type.BooleanType; import com.facebook.presto.spi.type.DateType; import com.facebook.presto.spi.type.DecimalType; import com.facebook.presto.spi.type.Decimals; import com.facebook.presto.spi.type.DoubleType; import com.facebook.presto.spi.type.IntegerType; import com.facebook.presto.spi.type.NamedTypeSignature; import com.facebook.presto.spi.type.SmallintType; import com.facebook.presto.spi.type.TimeType; import com.facebook.presto.spi.type.TimestampType; import com.facebook.presto.spi.type.TimestampWithTimeZoneType; import com.facebook.presto.spi.type.TinyintType; import com.facebook.presto.spi.type.Type; import com.facebook.presto.spi.type.TypeSignatureParameter; import com.facebook.presto.spi.type.VarbinaryType; import com.google.common.collect.ImmutableList; import com.mongodb.client.MongoCollection; import com.mongodb.client.model.InsertManyOptions; import io.airlift.slice.Slice; import org.bson.Document; import org.bson.types.Binary; import org.bson.types.ObjectId; import java.math.BigDecimal; import java.math.BigInteger; import java.util.ArrayList; import java.util.Collection; import java.util.Date; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; import static com.facebook.presto.mongodb.ObjectIdType.OBJECT_ID; import static com.facebook.presto.mongodb.TypeUtils.isArrayType; import static com.facebook.presto.mongodb.TypeUtils.isMapType; import static com.facebook.presto.mongodb.TypeUtils.isRowType; import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; import static com.facebook.presto.spi.type.DateTimeEncoding.unpackMillisUtc; import static com.facebook.presto.spi.type.Varchars.isVarcharType; import static java.util.Collections.unmodifiableList; import static java.util.Collections.unmodifiableMap; import static java.util.concurrent.CompletableFuture.completedFuture; public class MongoPageSink implements ConnectorPageSink { private final MongoSession mongoSession; private final ConnectorSession session; private final SchemaTableName schemaTableName; private final List<MongoColumnHandle> columns; private final String implicitPrefix; public MongoPageSink(MongoClientConfig config, MongoSession mongoSession, ConnectorSession session, SchemaTableName schemaTableName, List<MongoColumnHandle> columns) { this.mongoSession = mongoSession; this.session = session; this.schemaTableName = schemaTableName; this.columns = columns; this.implicitPrefix = config.getImplicitRowFieldPrefix(); } @Override public CompletableFuture<?> appendPage(Page page) { MongoCollection<Document> collection = mongoSession.getCollection(schemaTableName); List<Document> batch = new ArrayList<>(page.getPositionCount()); for (int position = 0; position < page.getPositionCount(); position++) { Document doc = new Document(); for (int channel = 0; channel < page.getChannelCount(); channel++) { MongoColumnHandle column = columns.get(channel); doc.append(column.getName(), getObjectValue(columns.get(channel).getType(), page.getBlock(channel), position)); } batch.add(doc); } collection.insertMany(batch, new InsertManyOptions().ordered(true)); return NOT_BLOCKED; } private Object getObjectValue(Type type, Block block, int position) { if (block.isNull(position)) { if (type.equals(OBJECT_ID)) { return new ObjectId(); } return null; } if (type.equals(OBJECT_ID)) { return new ObjectId(block.getSlice(position, 0, block.getLength(position)).getBytes()); } if (type.equals(BooleanType.BOOLEAN)) { return type.getBoolean(block, position); } if (type.equals(BigintType.BIGINT)) { return type.getLong(block, position); } if (type.equals(IntegerType.INTEGER)) { return (int) type.getLong(block, position); } if (type.equals(SmallintType.SMALLINT)) { return (short) type.getLong(block, position); } if (type.equals(TinyintType.TINYINT)) { return (byte) type.getLong(block, position); } if (type.equals(DoubleType.DOUBLE)) { return type.getDouble(block, position); } if (isVarcharType(type)) { return type.getSlice(block, position).toStringUtf8(); } if (type.equals(VarbinaryType.VARBINARY)) { return new Binary(type.getSlice(block, position).getBytes()); } if (type.equals(DateType.DATE)) { long days = type.getLong(block, position); return new Date(TimeUnit.DAYS.toMillis(days)); } if (type.equals(TimeType.TIME)) { long millisUtc = type.getLong(block, position); return new Date(millisUtc); } if (type.equals(TimestampType.TIMESTAMP)) { long millisUtc = type.getLong(block, position); return new Date(millisUtc); } if (type.equals(TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE)) { long millisUtc = unpackMillisUtc(type.getLong(block, position)); return new Date(millisUtc); } if (type instanceof DecimalType) { // TODO: decimal type might not support yet DecimalType decimalType = (DecimalType) type; BigInteger unscaledValue; if (decimalType.isShort()) { unscaledValue = BigInteger.valueOf(decimalType.getLong(block, position)); } else { unscaledValue = Decimals.decodeUnscaledValue(decimalType.getSlice(block, position)); } return new BigDecimal(unscaledValue); } if (isArrayType(type)) { Type elementType = type.getTypeParameters().get(0); Block arrayBlock = block.getObject(position, Block.class); List<Object> list = new ArrayList<>(arrayBlock.getPositionCount()); for (int i = 0; i < arrayBlock.getPositionCount(); i++) { Object element = getObjectValue(elementType, arrayBlock, i); list.add(element); } return unmodifiableList(list); } if (isMapType(type)) { Type keyType = type.getTypeParameters().get(0); Type valueType = type.getTypeParameters().get(1); Block mapBlock = block.getObject(position, Block.class); // map type is converted into list of fixed keys document List<Object> values = new ArrayList<>(mapBlock.getPositionCount() / 2); for (int i = 0; i < mapBlock.getPositionCount(); i += 2) { Map<String, Object> mapValue = new HashMap<>(); mapValue.put("key", getObjectValue(keyType, mapBlock, i)); mapValue.put("value", getObjectValue(valueType, mapBlock, i + 1)); values.add(mapValue); } return unmodifiableList(values); } if (isRowType(type)) { Block rowBlock = block.getObject(position, Block.class); List<Type> fieldTypes = type.getTypeParameters(); if (fieldTypes.size() != rowBlock.getPositionCount()) { throw new PrestoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, "Expected row value field count does not match type field count"); } if (isImplicitRowType(type)) { List<Object> rowValue = new ArrayList<>(); for (int i = 0; i < rowBlock.getPositionCount(); i++) { Object element = getObjectValue(fieldTypes.get(i), rowBlock, i); rowValue.add(element); } return unmodifiableList(rowValue); } Map<String, Object> rowValue = new HashMap<>(); for (int i = 0; i < rowBlock.getPositionCount(); i++) { rowValue.put(type.getTypeSignature().getParameters().get(i).getNamedTypeSignature().getName(), getObjectValue(fieldTypes.get(i), rowBlock, i)); } return unmodifiableMap(rowValue); } throw new PrestoException(NOT_SUPPORTED, "unsupported type: " + type); } private boolean isImplicitRowType(Type type) { return type.getTypeSignature().getParameters().stream().map(TypeSignatureParameter::getNamedTypeSignature) .map(NamedTypeSignature::getName).allMatch(name -> name.startsWith(implicitPrefix)); } @Override public CompletableFuture<Collection<Slice>> finish() { return completedFuture(ImmutableList.of()); } @Override public void abort() { } }