Support enforcement of NOT NULL column declarations

This commit enforces NOT NULL column declarations on write
in the Presto engine, so it applies to all connectors.  The
existing Postgres and Mysql tests named testInsertIntoNotNullColumn
were changed to check for the new error message, and a new test
with the same name was added to TestIcebergSmoke.

One possible concern with this commit is that the error message
issued by the Presto engine when writing a null to a NOT NULL
column is a different message than the Connector might issue
if no value was supplied for the NOT NULL column.  I think this
is ok, because the error messages supplied by the Connectors are
completely specific to the Connector.

Author: David Stryker

presto-iceberg/src/main/java/io/prestosql/plugin/iceberg/IcebergMetadata.java

@@ -543,6 +543,7 @@ private List<ColumnMetadata> getColumnMetadatas(org.apache.iceberg.Table table)
                     return ColumnMetadata.builder()
                             .setName(column.name())
                             .setType(toPrestoType(column.type(), typeManager))
+                            .setNullable(column.isOptional())
                             .setComment(Optional.ofNullable(column.doc()))
                             .build();
                 })

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

@@ -391,6 +391,22 @@ public void testSchemaEvolution()
         testWithAllFileFormats(this::testSchemaEvolution);
     }
 
+    @Test
+    public void testInsertIntoNotNullColumn()
+    {
+        assertUpdate("CREATE TABLE test_not_null_table (c1 INTEGER, c2 INTEGER NOT NULL)");
+        assertUpdate("INSERT INTO test_not_null_table (c2) VALUES (2)", 1);
+        assertQuery("SELECT * FROM test_not_null_table", "VALUES (NULL, 2)");
+        assertQueryFails("INSERT INTO test_not_null_table (c1) VALUES (1)", "NULL value not allowed for NOT NULL column: c2");
+        assertUpdate("DROP TABLE IF EXISTS test_not_null_table");
+
+        assertUpdate("CREATE TABLE test_commuted_not_null_table (a BIGINT, b BIGINT NOT NULL)");
+        assertUpdate("INSERT INTO test_commuted_not_null_table (b) VALUES (2)", 1);
+        assertQuery("SELECT * FROM test_commuted_not_null_table", "VALUES (NULL, 2)");
+        assertQueryFails("INSERT INTO test_commuted_not_null_table (b, a) VALUES (NULL, 3)", "NULL value not allowed for NOT NULL column: b");
+        assertUpdate("DROP TABLE IF EXISTS test_commuted_not_null_table");
+    }
+
     private void testSchemaEvolution(Session session, FileFormat fileFormat)
     {
         assertUpdate(session, "CREATE TABLE test_schema_evolution_drop_end (col0 INTEGER, col1 INTEGER, col2 INTEGER) WITH (format = '" + fileFormat + "')");

presto-main/src/main/java/io/prestosql/operator/TableWriterOperator.java

@@ -25,6 +25,7 @@
 import io.prestosql.operator.OperationTimer.OperationTiming;
 import io.prestosql.spi.Page;
 import io.prestosql.spi.PageBuilder;
+import io.prestosql.spi.PrestoException;
 import io.prestosql.spi.block.Block;
 import io.prestosql.spi.block.BlockBuilder;
 import io.prestosql.spi.block.RunLengthEncodedBlock;
@@ -48,6 +49,7 @@
 import static io.airlift.concurrent.MoreFutures.getFutureValue;
 import static io.airlift.concurrent.MoreFutures.toListenableFuture;
 import static io.prestosql.SystemSessionProperties.isStatisticsCpuTimerEnabled;
+import static io.prestosql.spi.StandardErrorCode.CONSTRAINT_VIOLATION;
 import static io.prestosql.spi.type.BigintType.BIGINT;
 import static io.prestosql.spi.type.VarbinaryType.VARBINARY;
 import static io.prestosql.sql.planner.plan.TableWriterNode.CreateTarget;
@@ -70,6 +72,7 @@ public static class TableWriterOperatorFactory
         private final PageSinkManager pageSinkManager;
         private final WriterTarget target;
         private final List<Integer> columnChannels;
+        private final List<String> notNullChannelColumnNames;
         private final Session session;
         private final OperatorFactory statisticsAggregationOperatorFactory;
         private final List<Type> types;
@@ -81,13 +84,15 @@ public TableWriterOperatorFactory(
                 PageSinkManager pageSinkManager,
                 WriterTarget writerTarget,
                 List<Integer> columnChannels,
+                List<String> notNullChannelColumnNames,
                 Session session,
                 OperatorFactory statisticsAggregationOperatorFactory,
                 List<Type> types)
         {
             this.operatorId = operatorId;
             this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
             this.columnChannels = requireNonNull(columnChannels, "columnChannels is null");
+            this.notNullChannelColumnNames = requireNonNull(notNullChannelColumnNames, "notNullChannelColumnNames is null");
             this.pageSinkManager = requireNonNull(pageSinkManager, "pageSinkManager is null");
             checkArgument(writerTarget instanceof CreateTarget || writerTarget instanceof InsertTarget, "writerTarget must be CreateTarget or InsertTarget");
             this.target = requireNonNull(writerTarget, "writerTarget is null");
@@ -103,7 +108,7 @@ public Operator createOperator(DriverContext driverContext)
             OperatorContext context = driverContext.addOperatorContext(operatorId, planNodeId, TableWriterOperator.class.getSimpleName());
             Operator statisticsAggregationOperator = statisticsAggregationOperatorFactory.createOperator(driverContext);
             boolean statisticsCpuTimerEnabled = !(statisticsAggregationOperator instanceof DevNullOperator) && isStatisticsCpuTimerEnabled(session);
-            return new TableWriterOperator(context, createPageSink(), columnChannels, statisticsAggregationOperator, types, statisticsCpuTimerEnabled);
+            return new TableWriterOperator(context, createPageSink(), columnChannels, notNullChannelColumnNames, statisticsAggregationOperator, types, statisticsCpuTimerEnabled);
         }
 
         private ConnectorPageSink createPageSink()
@@ -126,7 +131,7 @@ public void noMoreOperators()
         @Override
         public OperatorFactory duplicate()
         {
-            return new TableWriterOperatorFactory(operatorId, planNodeId, pageSinkManager, target, columnChannels, session, statisticsAggregationOperatorFactory, types);
+            return new TableWriterOperatorFactory(operatorId, planNodeId, pageSinkManager, target, columnChannels, notNullChannelColumnNames, session, statisticsAggregationOperatorFactory, types);
         }
     }
 
@@ -139,6 +144,7 @@ private enum State
     private final LocalMemoryContext pageSinkMemoryContext;
     private final ConnectorPageSink pageSink;
     private final List<Integer> columnChannels;
+    private final List<String> notNullChannelColumnNames;
     private final AtomicLong pageSinkPeakMemoryUsage = new AtomicLong();
     private final Operator statisticAggregationOperator;
     private final List<Type> types;
@@ -158,6 +164,7 @@ public TableWriterOperator(
             OperatorContext operatorContext,
             ConnectorPageSink pageSink,
             List<Integer> columnChannels,
+            List<String> notNullChannelColumnNames,
             Operator statisticAggregationOperator,
             List<Type> types,
             boolean statisticsCpuTimerEnabled)
@@ -166,6 +173,8 @@ public TableWriterOperator(
         this.pageSinkMemoryContext = operatorContext.newLocalSystemMemoryContext(TableWriterOperator.class.getSimpleName());
         this.pageSink = requireNonNull(pageSink, "pageSink is null");
         this.columnChannels = requireNonNull(columnChannels, "columnChannels is null");
+        this.notNullChannelColumnNames = requireNonNull(notNullChannelColumnNames, "notNullChannelColumnNames is null");
+        checkArgument(columnChannels.size() == notNullChannelColumnNames.size(), "columnChannels and notNullColumnNames have different sizes");
         this.operatorContext.setInfoSupplier(this::getInfo);
         this.statisticAggregationOperator = requireNonNull(statisticAggregationOperator, "statisticAggregationOperator is null");
         this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
@@ -227,7 +236,12 @@ public void addInput(Page page)
 
         Block[] blocks = new Block[columnChannels.size()];
         for (int outputChannel = 0; outputChannel < columnChannels.size(); outputChannel++) {
-            blocks[outputChannel] = page.getBlock(columnChannels.get(outputChannel));
+            Block block = page.getBlock(columnChannels.get(outputChannel));
+            String columnName = notNullChannelColumnNames.get(outputChannel);
+            if (columnName != null) {
+                verifyBlockHasNoNulls(block, columnName);
+            }
+            blocks[outputChannel] = block;
         }
 
         OperationTimer timer = new OperationTimer(statisticsCpuTimerEnabled);
@@ -243,6 +257,18 @@ public void addInput(Page page)
         updateWrittenBytes();
     }
 
+    private void verifyBlockHasNoNulls(Block block, String columnName)
+    {
+        if (!block.mayHaveNull()) {
+            return;
+        }
+        for (int position = 0; position < block.getPositionCount(); position++) {
+            if (block.isNull(position)) {
+                throw new PrestoException(CONSTRAINT_VIOLATION, "NULL value not allowed for NOT NULL column: " + columnName);
+            }
+        }
+    }
+
     @Override
     public Page getOutput()
     {

presto-main/src/main/java/io/prestosql/sql/planner/LocalExecutionPlanner.java

@@ -2335,12 +2335,17 @@ public PhysicalOperation visitTableWriter(TableWriterNode node, LocalExecutionPl
                     .map(source::symbolToChannel)
                     .collect(toImmutableList());
 
+            List<String> notNullChannelColumnNames = node.getColumns().stream()
+                    .map(symbol -> node.getNotNullColumnSymbols().contains(symbol) ? node.getColumnNames().get(source.symbolToChannel(symbol)) : null)
+                    .collect(Collectors.toList());
+
             OperatorFactory operatorFactory = new TableWriterOperatorFactory(
                     context.getNextOperatorId(),
                     node.getId(),
                     pageSinkManager,
                     node.getTarget(),
                     inputChannels,
+                    notNullChannelColumnNames,
                     session,
                     statisticsAggregation,
                     getSymbolTypes(node.getOutputSymbols(), context.getTypes()));

presto-main/src/main/java/io/prestosql/sql/planner/LogicalPlanner.java

@@ -92,11 +92,13 @@
 import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Streams.zip;
 import static io.prestosql.SystemSessionProperties.isCollectPlanStatisticsForAllQueries;
 import static io.prestosql.SystemSessionProperties.isUsePreferredWritePartitioning;
@@ -334,6 +336,7 @@ private RelationPlan createTableCreationPlan(Analysis analysis, Query query)
                 plan,
                 new CreateReference(destination.getCatalogName(), tableMetadata, newTableLayout),
                 columnNames,
+                tableMetadata.getColumns(),
                 newTableLayout,
                 statisticsMetadata);
     }
@@ -407,6 +410,7 @@ private RelationPlan createInsertPlan(Analysis analysis, Insert insertStatement)
                                 .map(columns::get)
                                 .collect(toImmutableList())),
                 insertedTableColumnNames,
+                insertedColumns,
                 insert.getNewTableLayout(),
                 statisticsMetadata);
     }
@@ -416,6 +420,7 @@ private RelationPlan createTableWriterPlan(
             RelationPlan plan,
             WriterTarget target,
             List<String> columnNames,
+            List<ColumnMetadata> columnMetadataList,
             Optional<NewTableLayout> writeTableLayout,
             TableStatisticsMetadata statisticsMetadata)
     {
@@ -448,11 +453,17 @@ else if (isUsePreferredWritePartitioning(session)) {
             }
         }
 
-        if (!statisticsMetadata.isEmpty()) {
-            verify(columnNames.size() == symbols.size(), "columnNames.size() != symbols.size(): %s and %s", columnNames, symbols);
-            Map<String, Symbol> columnToSymbolMap = zip(columnNames.stream(), symbols.stream(), SimpleImmutableEntry::new)
-                    .collect(toImmutableMap(Entry::getKey, Entry::getValue));
+        verify(columnNames.size() == symbols.size(), "columnNames.size() != symbols.size(): %s and %s", columnNames, symbols);
+        Map<String, Symbol> columnToSymbolMap = zip(columnNames.stream(), symbols.stream(), SimpleImmutableEntry::new)
+                .collect(toImmutableMap(Entry::getKey, Entry::getValue));
+
+        Set<Symbol> notNullColumnSymbols = columnMetadataList.stream()
+                .filter(column -> !column.isNullable())
+                .map(ColumnMetadata::getName)
+                .map(columnToSymbolMap::get)
+                .collect(toImmutableSet());
 
+        if (!statisticsMetadata.isEmpty()) {
             TableStatisticAggregation result = statisticsAggregationPlanner.createStatisticsAggregation(statisticsMetadata, columnToSymbolMap);
 
             StatisticAggregations.Parts aggregations = result.getAggregations().createPartialAggregations(symbolAllocator, metadata);
@@ -473,6 +484,7 @@ else if (isUsePreferredWritePartitioning(session)) {
                             symbolAllocator.newSymbol("fragment", VARBINARY),
                             symbols,
                             columnNames,
+                            notNullColumnSymbols,
                             partitioningScheme,
                             Optional.of(partialAggregation),
                             Optional.of(result.getDescriptor().map(aggregations.getMappings()::get))),
@@ -494,6 +506,7 @@ else if (isUsePreferredWritePartitioning(session)) {
                         symbolAllocator.newSymbol("fragment", VARBINARY),
                         symbols,
                         columnNames,
+                        notNullColumnSymbols,
                         partitioningScheme,
                         Optional.empty(),
                         Optional.empty()),

presto-main/src/main/java/io/prestosql/sql/planner/optimizations/BeginTableWrite.java

@@ -98,6 +98,7 @@ public PlanNode visitTableWriter(TableWriterNode node, RewriteContext<Context> c
                     node.getFragmentSymbol(),
                     node.getColumns(),
                     node.getColumnNames(),
+                    node.getNotNullColumnSymbols(),
                     node.getPartitioningScheme(),
                     node.getStatisticsAggregation(),
                     node.getStatisticsAggregationDescriptor());

presto-main/src/main/java/io/prestosql/sql/planner/optimizations/PruneUnreferencedOutputs.java

@@ -681,6 +681,7 @@ public PlanNode visitTableWriter(TableWriterNode node, RewriteContext<Set<Symbol
                     node.getFragmentSymbol(),
                     node.getColumns(),
                     node.getColumnNames(),
+                    node.getNotNullColumnSymbols(),
                     node.getPartitioningScheme(),
                     node.getStatisticsAggregation(),
                     node.getStatisticsAggregationDescriptor());

presto-main/src/main/java/io/prestosql/sql/planner/optimizations/SymbolMapper.java

@@ -162,6 +162,7 @@ public TableWriterNode map(TableWriterNode node, PlanNode source, PlanNodeId new
                 map(node.getFragmentSymbol()),
                 columns,
                 node.getColumnNames(),
+                node.getNotNullColumnSymbols(),
                 node.getPartitioningScheme().map(partitioningScheme -> canonicalize(partitioningScheme, source)),
                 node.getStatisticsAggregation().map(this::map),
                 node.getStatisticsAggregationDescriptor().map(this::map));

presto-main/src/main/java/io/prestosql/sql/planner/plan/TableWriterNode.java

@@ -18,6 +18,7 @@
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import io.prestosql.metadata.InsertTableHandle;
 import io.prestosql.metadata.NewTableLayout;
@@ -33,6 +34,7 @@
 
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
@@ -47,6 +49,7 @@ public class TableWriterNode
     private final Symbol fragmentSymbol;
     private final List<Symbol> columns;
     private final List<String> columnNames;
+    private final Set<Symbol> notNullColumnSymbols;
     private final Optional<PartitioningScheme> partitioningScheme;
     private final Optional<StatisticAggregations> statisticsAggregation;
     private final Optional<StatisticAggregationsDescriptor<Symbol>> statisticsAggregationDescriptor;
@@ -61,6 +64,7 @@ public TableWriterNode(
             @JsonProperty("fragmentSymbol") Symbol fragmentSymbol,
             @JsonProperty("columns") List<Symbol> columns,
             @JsonProperty("columnNames") List<String> columnNames,
+            @JsonProperty("notNullColumnSymbols") Set<Symbol> notNullColumnSymbols,
             @JsonProperty("partitioningScheme") Optional<PartitioningScheme> partitioningScheme,
             @JsonProperty("statisticsAggregation") Optional<StatisticAggregations> statisticsAggregation,
             @JsonProperty("statisticsAggregationDescriptor") Optional<StatisticAggregationsDescriptor<Symbol>> statisticsAggregationDescriptor)
@@ -77,6 +81,7 @@ public TableWriterNode(
         this.fragmentSymbol = requireNonNull(fragmentSymbol, "fragmentSymbol is null");
         this.columns = ImmutableList.copyOf(columns);
         this.columnNames = ImmutableList.copyOf(columnNames);
+        this.notNullColumnSymbols = ImmutableSet.copyOf(requireNonNull(notNullColumnSymbols, "notNullColumns is null"));
         this.partitioningScheme = requireNonNull(partitioningScheme, "partitioningScheme is null");
         this.statisticsAggregation = requireNonNull(statisticsAggregation, "statisticsAggregation is null");
         this.statisticsAggregationDescriptor = requireNonNull(statisticsAggregationDescriptor, "statisticsAggregationDescriptor is null");
@@ -128,6 +133,12 @@ public List<String> getColumnNames()
         return columnNames;
     }
 
+    @JsonProperty
+    public Set<Symbol> getNotNullColumnSymbols()
+    {
+        return notNullColumnSymbols;
+    }
+
     @JsonProperty
     public Optional<PartitioningScheme> getPartitioningScheme()
     {
@@ -167,7 +178,7 @@ public <R, C> R accept(PlanVisitor<R, C> visitor, C context)
     @Override
     public PlanNode replaceChildren(List<PlanNode> newChildren)
     {
-        return new TableWriterNode(getId(), Iterables.getOnlyElement(newChildren), target, rowCountSymbol, fragmentSymbol, columns, columnNames, partitioningScheme, statisticsAggregation, statisticsAggregationDescriptor);
+        return new TableWriterNode(getId(), Iterables.getOnlyElement(newChildren), target, rowCountSymbol, fragmentSymbol, columns, columnNames, notNullColumnSymbols, partitioningScheme, statisticsAggregation, statisticsAggregationDescriptor);
     }
 
     @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "@type")

presto-main/src/test/java/io/prestosql/operator/TestTableWriterOperator.java

@@ -292,6 +292,8 @@ private Operator createTableWriterOperator(
             Session session,
             DriverContext driverContext)
     {
+        List<String> notNullColumnNames = new ArrayList<>(1);
+        notNullColumnNames.add(null);
         TableWriterOperatorFactory factory = new TableWriterOperatorFactory(
                 0,
                 new PlanNodeId("test"),
@@ -302,6 +304,7 @@ private Operator createTableWriterOperator(
                         new ConnectorOutputTableHandle() {}),
                         new SchemaTableName("testSchema", "testTable")),
                 ImmutableList.of(0),
+                notNullColumnNames,
                 session,
                 statisticsAggregation,
                 outputTypes);