Disable incorrect pushdown in Mongo connector

presto-accumulo/src/test/java/io/prestosql/plugin/accumulo/TestAccumuloDistributedQueries.java

@@ -23,6 +23,8 @@
 import org.testng.SkipException;
 import org.testng.annotations.Test;
 
+import java.util.Optional;
+
 import static io.prestosql.plugin.accumulo.AccumuloQueryRunner.createAccumuloQueryRunner;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
@@ -358,4 +360,17 @@ public void testCommentTable()
         // Accumulo connector currently does not support comment on table
         assertQueryFails("COMMENT ON TABLE orders IS 'hello'", "This connector does not support setting table comments");
     }
+
+    @Override
+    protected Optional<DataMappingTestSetup> filterDataMappingSmokeTestData(DataMappingTestSetup dataMappingTestSetup)
+    {
+        String typeName = dataMappingTestSetup.getPrestoTypeName();
+        if (typeName.startsWith("decimal(")
+                || typeName.equals("timestamp with time zone")
+                || typeName.startsWith("char(")) {
+            return Optional.of(dataMappingTestSetup.asUnsupported());
+        }
+
+        return Optional.of(dataMappingTestSetup);
+    }
 }

presto-cassandra/src/test/java/io/prestosql/plugin/cassandra/TestCassandraDistributedQueries.java

@@ -171,4 +171,10 @@ protected TestTable createTableWithDefaultColumns()
     {
         throw new SkipException("Cassandra connector does not support column default values");
     }
+
+    @Override
+    public void testDataMappingSmokeTest(DataMappingTestSetup dataMappingTestSetup)
+    {
+        // Cassandra connector currently does not support create table
+    }
 }

presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveDistributedQueries.java

@@ -20,6 +20,8 @@
 import org.testng.SkipException;
 import org.testng.annotations.Test;
 
+import java.util.Optional;
+
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static io.prestosql.sql.tree.ExplainType.Type.LOGICAL;
 import static io.prestosql.tpch.TpchTable.getTables;
@@ -57,5 +59,17 @@ public void testExplainOfCreateTableAs()
         assertEquals(getOnlyElement(result.getOnlyColumnAsSet()), getExplainPlan(query, LOGICAL));
     }
 
+    @Override
+    protected Optional<DataMappingTestSetup> filterDataMappingSmokeTestData(DataMappingTestSetup dataMappingTestSetup)
+    {
+        String typeName = dataMappingTestSetup.getPrestoTypeName();
+        if (typeName.equals("time")
+                || typeName.equals("timestamp with time zone")) {
+            return Optional.of(dataMappingTestSetup.asUnsupported());
+        }
+
+        return Optional.of(dataMappingTestSetup);
+    }
+
     // Hive specific tests should normally go in TestHiveIntegrationSmokeTest
 }

presto-iceberg/src/test/java/io/prestosql/plugin/iceberg/TestIcebergDistributed.java

@@ -19,6 +19,8 @@
 import io.prestosql.testing.sql.TestTable;
 import org.testng.SkipException;
 
+import java.util.Optional;
+
 import static io.prestosql.plugin.iceberg.IcebergQueryRunner.createIcebergQueryRunner;
 
 public class TestIcebergDistributed
@@ -79,4 +81,25 @@ public void testInsertWithCoercion()
     {
         // Iceberg does not support parameterized varchar
     }
+
+    @Override
+    protected Optional<DataMappingTestSetup> filterDataMappingSmokeTestData(DataMappingTestSetup dataMappingTestSetup)
+    {
+        String typeName = dataMappingTestSetup.getPrestoTypeName();
+        if (typeName.equals("tinyint")
+                || typeName.equals("smallint")
+                || typeName.equals("timestamp")
+                || typeName.startsWith("char(")) {
+            return Optional.of(dataMappingTestSetup.asUnsupported());
+        }
+
+        if (typeName.startsWith("decimal(")
+                || typeName.equals("time")
+                || typeName.equals("timestamp with time zone")) {
+            // TODO this should either work or fail cleanly
+            return Optional.empty();
+        }
+
+        return Optional.of(dataMappingTestSetup);
+    }
 }

presto-mongodb/src/main/java/io/prestosql/plugin/mongodb/MongoPageSink.java

@@ -28,12 +28,13 @@
 import io.prestosql.spi.connector.SchemaTableName;
 import io.prestosql.spi.type.BigintType;
 import io.prestosql.spi.type.BooleanType;
+import io.prestosql.spi.type.CharType;
 import io.prestosql.spi.type.DateType;
 import io.prestosql.spi.type.DecimalType;
-import io.prestosql.spi.type.Decimals;
 import io.prestosql.spi.type.DoubleType;
 import io.prestosql.spi.type.IntegerType;
 import io.prestosql.spi.type.NamedTypeSignature;
+import io.prestosql.spi.type.RealType;
 import io.prestosql.spi.type.SmallintType;
 import io.prestosql.spi.type.TimeType;
 import io.prestosql.spi.type.TimestampType;
@@ -46,8 +47,6 @@
 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;
@@ -63,8 +62,11 @@
 import static io.prestosql.plugin.mongodb.TypeUtils.isMapType;
 import static io.prestosql.plugin.mongodb.TypeUtils.isRowType;
 import static io.prestosql.spi.StandardErrorCode.NOT_SUPPORTED;
+import static io.prestosql.spi.type.Chars.padSpaces;
 import static io.prestosql.spi.type.DateTimeEncoding.unpackMillisUtc;
+import static io.prestosql.spi.type.Decimals.readBigDecimal;
 import static io.prestosql.spi.type.Varchars.isVarcharType;
+import static java.lang.Float.intBitsToFloat;
 import static java.lang.Math.toIntExact;
 import static java.util.Collections.unmodifiableList;
 import static java.util.Collections.unmodifiableMap;
@@ -141,12 +143,18 @@ private Object getObjectValue(Type type, Block block, int position)
         if (type.equals(TinyintType.TINYINT)) {
             return SignedBytes.checkedCast(type.getLong(block, position));
         }
+        if (type.equals(RealType.REAL)) {
+            return intBitsToFloat(toIntExact(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 instanceof CharType) {
+            return padSpaces(type.getSlice(block, position), ((CharType) type)).toStringUtf8();
+        }
         if (type.equals(VarbinaryType.VARBINARY)) {
             return new Binary(type.getSlice(block, position).getBytes());
         }
@@ -167,17 +175,7 @@ private Object getObjectValue(Type type, Block block, int position)
             return new Date(millisUtc);
         }
         if (type instanceof DecimalType) {
-            // TODO: decimal type might not support yet
-            // TODO: this code is likely wrong and should switch to Decimals.readBigDecimal()
-            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);
+            return readBigDecimal((DecimalType) type, block, position);
         }
         if (isArrayType(type)) {
             Type elementType = type.getTypeParameters().get(0);

presto-mongodb/src/main/java/io/prestosql/plugin/mongodb/MongoPageSource.java

@@ -13,6 +13,8 @@
  */
 package io.prestosql.plugin.mongodb;
 
+import com.google.common.primitives.Shorts;
+import com.google.common.primitives.SignedBytes;
 import com.mongodb.client.MongoCursor;
 import io.airlift.slice.Slice;
 import io.prestosql.spi.Page;
@@ -21,12 +23,15 @@
 import io.prestosql.spi.block.Block;
 import io.prestosql.spi.block.BlockBuilder;
 import io.prestosql.spi.connector.ConnectorPageSource;
+import io.prestosql.spi.type.CharType;
+import io.prestosql.spi.type.DecimalType;
 import io.prestosql.spi.type.Type;
 import io.prestosql.spi.type.TypeSignatureParameter;
 import io.prestosql.spi.type.VarbinaryType;
 import io.prestosql.spi.type.VarcharType;
 import org.bson.Document;
 import org.bson.types.Binary;
+import org.bson.types.Decimal128;
 import org.bson.types.ObjectId;
 import org.joda.time.chrono.ISOChronology;
 
@@ -47,10 +52,20 @@
 import static io.prestosql.plugin.mongodb.TypeUtils.isRowType;
 import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static io.prestosql.spi.type.BigintType.BIGINT;
+import static io.prestosql.spi.type.Chars.truncateToLengthAndTrimSpaces;
+import static io.prestosql.spi.type.DateTimeEncoding.packDateTimeWithZone;
 import static io.prestosql.spi.type.DateType.DATE;
+import static io.prestosql.spi.type.Decimals.encodeScaledValue;
+import static io.prestosql.spi.type.Decimals.encodeShortScaledValue;
 import static io.prestosql.spi.type.IntegerType.INTEGER;
+import static io.prestosql.spi.type.RealType.REAL;
+import static io.prestosql.spi.type.SmallintType.SMALLINT;
 import static io.prestosql.spi.type.TimeType.TIME;
+import static io.prestosql.spi.type.TimeZoneKey.UTC_KEY;
 import static io.prestosql.spi.type.TimestampType.TIMESTAMP;
+import static io.prestosql.spi.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE;
+import static io.prestosql.spi.type.TinyintType.TINYINT;
+import static java.lang.Float.floatToIntBits;
 import static java.lang.String.join;
 import static java.util.stream.Collectors.toList;
 
@@ -150,6 +165,19 @@ else if (javaType == long.class) {
                 else if (type.equals(INTEGER)) {
                     type.writeLong(output, ((Number) value).intValue());
                 }
+                else if (type.equals(SMALLINT)) {
+                    type.writeLong(output, Shorts.checkedCast(((Number) value).longValue()));
+                }
+                else if (type.equals(TINYINT)) {
+                    type.writeLong(output, SignedBytes.checkedCast(((Number) value).longValue()));
+                }
+                else if (type.equals(REAL)) {
+                    //noinspection NumericCastThatLosesPrecision
+                    type.writeLong(output, floatToIntBits(((float) ((Number) value).doubleValue())));
+                }
+                else if (type instanceof DecimalType) {
+                    type.writeLong(output, encodeShortScaledValue(((Decimal128) value).bigDecimalValue(), ((DecimalType) type).getScale()));
+                }
                 else if (type.equals(DATE)) {
                     long utcMillis = ((Date) value).getTime();
                     type.writeLong(output, TimeUnit.MILLISECONDS.toDays(utcMillis));
@@ -158,8 +186,12 @@ else if (type.equals(TIME)) {
                     type.writeLong(output, UTC_CHRONOLOGY.millisOfDay().get(((Date) value).getTime()));
                 }
                 else if (type.equals(TIMESTAMP)) {
+                    // TODO provide correct TIMESTAMP mapping, and respecting session.isLegacyTimestamp()
                     type.writeLong(output, ((Date) value).getTime());
                 }
+                else if (type.equals(TIMESTAMP_WITH_TIME_ZONE)) {
+                    type.writeLong(output, packDateTimeWithZone(((Date) value).getTime(), UTC_KEY));
+                }
                 else {
                     throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unhandled type for " + javaType.getSimpleName() + ":" + type.getTypeSignature());
                 }
@@ -178,6 +210,7 @@ else if (javaType == Block.class) {
             }
         }
         catch (ClassCastException ignore) {
+            // TODO remove (fail clearly), or hide behind a toggle
             // returns null instead of raising exception
             output.appendNull();
         }
@@ -199,6 +232,9 @@ private void writeSlice(BlockBuilder output, Type type, Object value)
         if (type instanceof VarcharType) {
             type.writeSlice(output, utf8Slice(toVarcharValue(value)));
         }
+        else if (type instanceof CharType) {
+            type.writeSlice(output, truncateToLengthAndTrimSpaces(utf8Slice((String) value), ((CharType) type)));
+        }
         else if (type.equals(OBJECT_ID)) {
             type.writeSlice(output, wrappedBuffer(((ObjectId) value).toByteArray()));
         }
@@ -210,6 +246,9 @@ else if (type instanceof VarbinaryType) {
                 output.appendNull();
             }
         }
+        else if (type instanceof DecimalType) {
+            type.writeSlice(output, encodeScaledValue(((Decimal128) value).bigDecimalValue(), ((DecimalType) type).getScale()));
+        }
         else {
             throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unhandled type for Slice: " + type.getTypeSignature());
         }