feat: Support vector search in LogicalPlanner (#27169)

presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java

@@ -186,6 +186,9 @@ public class Analysis
     private Optional<TableHandle> callTarget = Optional.empty();
     private Optional<QuerySpecification> targetQuery = Optional.empty();
 
+    // for create vector index
+    private Optional<CreateVectorIndexAnalysis> createVectorIndexAnalysis = Optional.empty();
+
     // for create table
     private Optional<QualifiedObjectName> createTableDestination = Optional.empty();
     private Map<String, Expression> createTableProperties = ImmutableMap.of();
@@ -700,6 +703,16 @@ public Optional<QualifiedObjectName> getCreateTableDestination()
         return createTableDestination;
     }
 
+    public void setCreateVectorIndexAnalysis(CreateVectorIndexAnalysis analysis)
+    {
+        this.createVectorIndexAnalysis = Optional.of(analysis);
+    }
+
+    public Optional<CreateVectorIndexAnalysis> getCreateVectorIndexAnalysis()
+    {
+        return createVectorIndexAnalysis;
+    }
+
     public Optional<QualifiedObjectName> getProcedureName()
     {
         return procedureName;
@@ -1937,4 +1950,53 @@ public Scope getTargetTableScope()
             return targetTableScope;
         }
     }
+
+    @Immutable
+    public static final class CreateVectorIndexAnalysis
+    {
+        private final QualifiedObjectName sourceTableName;
+        private final QualifiedObjectName targetTableName;
+        private final List<Identifier> columns;
+        private final Map<String, Expression> properties;
+        private final Optional<Expression> updatingFor;
+
+        public CreateVectorIndexAnalysis(
+                QualifiedObjectName sourceTableName,
+                QualifiedObjectName targetTableName,
+                List<Identifier> columns,
+                Map<String, Expression> properties,
+                Optional<Expression> updatingFor)
+        {
+            this.sourceTableName = requireNonNull(sourceTableName, "sourceTableName is null");
+            this.targetTableName = requireNonNull(targetTableName, "targetTableName is null");
+            this.columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null"));
+            this.properties = ImmutableMap.copyOf(requireNonNull(properties, "properties is null"));
+            this.updatingFor = requireNonNull(updatingFor, "updatingFor is null");
+        }
+
+        public QualifiedObjectName getSourceTableName()
+        {
+            return sourceTableName;
+        }
+
+        public QualifiedObjectName getTargetTableName()
+        {
+            return targetTableName;
+        }
+
+        public List<Identifier> getColumns()
+        {
+            return columns;
+        }
+
+        public Map<String, Expression> getProperties()
+        {
+            return properties;
+        }
+
+        public Optional<Expression> getUpdatingFor()
+        {
+            return updatingFor;
+        }
+    }
 }

presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/utils/StatementUtils.java

@@ -30,6 +30,7 @@
 import com.facebook.presto.sql.tree.CreateTableAsSelect;
 import com.facebook.presto.sql.tree.CreateTag;
 import com.facebook.presto.sql.tree.CreateType;
+import com.facebook.presto.sql.tree.CreateVectorIndex;
 import com.facebook.presto.sql.tree.CreateView;
 import com.facebook.presto.sql.tree.Deallocate;
 import com.facebook.presto.sql.tree.Delete;
@@ -107,6 +108,7 @@ private StatementUtils() {}
         builder.put(CreateTableAsSelect.class, QueryType.INSERT);
         builder.put(Insert.class, QueryType.INSERT);
         builder.put(RefreshMaterializedView.class, QueryType.INSERT);
+        builder.put(CreateVectorIndex.class, QueryType.INSERT);
 
         builder.put(Delete.class, QueryType.DELETE);
         builder.put(Update.class, QueryType.UPDATE);

presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/TableWriteInfo.java

@@ -22,6 +22,7 @@
 import com.facebook.presto.metadata.AnalyzeTableHandle;
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.spi.MergeHandle;
+import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.plan.PlanNode;
 import com.facebook.presto.spi.plan.TableFinishNode;
 import com.facebook.presto.spi.plan.TableWriterNode;
@@ -36,6 +37,7 @@
 import java.util.List;
 import java.util.Optional;
 
+import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Iterables.getOnlyElement;
@@ -122,6 +124,11 @@ private static Optional<ExecutionWriterTarget> createWriterTarget(Optional<Table
                 return Optional.of(new ExecutionWriterTarget.MergeHandle(mergeHandle.orElseThrow(
                         () -> new VerifyException("mergeHandle is absent: " + target.getClass().getSimpleName()))));
             }
+            if (target instanceof TableWriterNode.CreateVectorIndexReference) {
+                throw new PrestoException(NOT_SUPPORTED,
+                        "This connector does not support creating vector indexes. " +
+                        "The connector must provide a ConnectorPlanOptimizer to handle CREATE VECTOR INDEX.");
+            }
             throw new IllegalArgumentException("Unhandled target type: " + target.getClass().getSimpleName());
         }
 

presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java

@@ -62,6 +62,7 @@
 import com.facebook.presto.operator.aggregation.ClassificationThresholdsAggregation;
 import com.facebook.presto.operator.aggregation.CountAggregation;
 import com.facebook.presto.operator.aggregation.CountIfAggregation;
+import com.facebook.presto.operator.aggregation.CreateVectorIndexAggregation;
 import com.facebook.presto.operator.aggregation.DefaultApproximateCountDistinctAggregation;
 import com.facebook.presto.operator.aggregation.DoubleCorrelationAggregation;
 import com.facebook.presto.operator.aggregation.DoubleCovarianceAggregation;
@@ -710,6 +711,7 @@ private List<? extends SqlFunction> getBuiltInFunctions(FunctionsConfig function
                 .aggregate(GeometryUnionAgg.class)
                 .aggregate(SpatialPartitioningAggregateFunction.class)
                 .aggregate(SpatialPartitioningInternalAggregateFunction.class)
+                .aggregates(CreateVectorIndexAggregation.class)
                 .aggregates(CountAggregation.class)
                 .aggregates(VarianceAggregation.class)
                 .aggregates(CentralMomentsAggregation.class)

presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java

@@ -29,6 +29,7 @@
 import com.facebook.presto.spi.MaterializedViewStatus;
 import com.facebook.presto.spi.MergeHandle;
 import com.facebook.presto.spi.NewTableLayout;
+import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.SystemTable;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.TableMetadata;
@@ -310,6 +311,22 @@ public Optional<ConnectorOutputMetadata> finishCreateTable(
         return delegate.finishCreateTable(session, tableHandle, fragments, computedStatistics);
     }
 
+    @Override
+    public OutputTableHandle beginCreateVectorIndex(Session session, String catalogName, ConnectorTableMetadata indexMetadata, Optional<NewTableLayout> layout, SchemaTableName sourceTableName)
+    {
+        return delegate.beginCreateVectorIndex(session, catalogName, indexMetadata, layout, sourceTableName);
+    }
+
+    @Override
+    public Optional<ConnectorOutputMetadata> finishCreateVectorIndex(
+            Session session,
+            OutputTableHandle tableHandle,
+            Collection<Slice> fragments,
+            Collection<ComputedStatistics> computedStatistics)
+    {
+        return delegate.finishCreateVectorIndex(session, tableHandle, fragments, computedStatistics);
+    }
+
     @Override
     public Optional<NewTableLayout> getInsertLayout(Session session, TableHandle target)
     {

presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java

@@ -32,6 +32,7 @@
 import com.facebook.presto.spi.MergeHandle;
 import com.facebook.presto.spi.NewTableLayout;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.SystemTable;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.TableLayoutFilterCoverage;
@@ -69,6 +70,7 @@
 import java.util.OptionalLong;
 import java.util.Set;
 
+import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
 import static com.facebook.presto.spi.TableLayoutFilterCoverage.NOT_APPLICABLE;
 
 public interface Metadata
@@ -270,6 +272,22 @@ public interface Metadata
      */
     Optional<ConnectorOutputMetadata> finishCreateTable(Session session, OutputTableHandle tableHandle, Collection<Slice> fragments, Collection<ComputedStatistics> computedStatistics);
 
+    /**
+     * Begin the atomic creation of a vector index with data.
+     */
+    default OutputTableHandle beginCreateVectorIndex(Session session, String catalogName, ConnectorTableMetadata indexMetadata, Optional<NewTableLayout> layout, SchemaTableName sourceTableName)
+    {
+        throw new PrestoException(NOT_SUPPORTED, "This connector does not support creating vector indexes");
+    }
+
+    /**
+     * Finish a vector index creation with data after the data is written.
+     */
+    default Optional<ConnectorOutputMetadata> finishCreateVectorIndex(Session session, OutputTableHandle tableHandle, Collection<Slice> fragments, Collection<ComputedStatistics> computedStatistics)
+    {
+        throw new PrestoException(NOT_SUPPORTED, "This connector does not support creating vector indexes");
+    }
+
     Optional<NewTableLayout> getInsertLayout(Session session, TableHandle target);
 
     /**

presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java

@@ -957,6 +957,27 @@ public Optional<ConnectorOutputMetadata> finishCreateTable(Session session, Outp
         return metadata.finishCreateTable(session.toConnectorSession(connectorId), tableHandle.getConnectorHandle(), fragments, computedStatistics);
     }
 
+    @Override
+    public OutputTableHandle beginCreateVectorIndex(Session session, String catalogName, ConnectorTableMetadata indexMetadata, Optional<NewTableLayout> layout, SchemaTableName sourceTableName)
+    {
+        CatalogMetadata catalogMetadata = getCatalogMetadataForWrite(session, catalogName);
+        ConnectorId connectorId = catalogMetadata.getConnectorId();
+        ConnectorMetadata metadata = catalogMetadata.getMetadata();
+
+        ConnectorTransactionHandle transactionHandle = catalogMetadata.getTransactionHandleFor(connectorId);
+        ConnectorSession connectorSession = session.toConnectorSession(connectorId);
+        ConnectorOutputTableHandle handle = metadata.beginCreateVectorIndex(connectorSession, indexMetadata, layout.map(NewTableLayout::getLayout), sourceTableName);
+        return new OutputTableHandle(connectorId, transactionHandle, handle);
+    }
+
+    @Override
+    public Optional<ConnectorOutputMetadata> finishCreateVectorIndex(Session session, OutputTableHandle tableHandle, Collection<Slice> fragments, Collection<ComputedStatistics> computedStatistics)
+    {
+        ConnectorId connectorId = tableHandle.getConnectorId();
+        ConnectorMetadata metadata = getMetadata(session, connectorId);
+        return metadata.finishCreateVectorIndex(session.toConnectorSession(connectorId), tableHandle.getConnectorHandle(), fragments, computedStatistics);
+    }
+
     @Override
     public InsertTableHandle beginInsert(Session session, TableHandle tableHandle)
     {

presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManagerStats.java

@@ -55,6 +55,8 @@ public class MetadataManagerStats
     private final AtomicLong getNewTableLayoutCalls = new AtomicLong();
     private final AtomicLong beginCreateTableCalls = new AtomicLong();
     private final AtomicLong finishCreateTableCalls = new AtomicLong();
+    private final AtomicLong beginCreateVectorIndexCalls = new AtomicLong();
+    private final AtomicLong finishCreateVectorIndexCalls = new AtomicLong();
     private final AtomicLong getInsertLayoutCalls = new AtomicLong();
     private final AtomicLong getStatisticsCollectionMetadataForWriteCalls = new AtomicLong();
     private final AtomicLong getStatisticsCollectionMetadataCalls = new AtomicLong();
@@ -165,6 +167,8 @@ public class MetadataManagerStats
     private final TimeStat getNewTableLayoutTime = new TimeStat(TimeUnit.NANOSECONDS);
     private final TimeStat beginCreateTableTime = new TimeStat(TimeUnit.NANOSECONDS);
     private final TimeStat finishCreateTableTime = new TimeStat(TimeUnit.NANOSECONDS);
+    private final TimeStat beginCreateVectorIndexTime = new TimeStat(TimeUnit.NANOSECONDS);
+    private final TimeStat finishCreateVectorIndexTime = new TimeStat(TimeUnit.NANOSECONDS);
     private final TimeStat getInsertLayoutTime = new TimeStat(TimeUnit.NANOSECONDS);
     private final TimeStat getStatisticsCollectionMetadataForWriteTime = new TimeStat(TimeUnit.NANOSECONDS);
     private final TimeStat getStatisticsCollectionMetadataTime = new TimeStat(TimeUnit.NANOSECONDS);
@@ -648,6 +652,20 @@ public TimeStat getFinishCreateTableTime()
         return finishCreateTableTime;
     }
 
+    @Managed
+    @Nested
+    public TimeStat getBeginCreateVectorIndexTime()
+    {
+        return beginCreateVectorIndexTime;
+    }
+
+    @Managed
+    @Nested
+    public TimeStat getFinishCreateVectorIndexTime()
+    {
+        return finishCreateVectorIndexTime;
+    }
+
     @Managed
     @Nested
     public TimeStat getGetInsertLayoutTime()
@@ -1357,6 +1375,18 @@ public void recordFinishCreateTableCall(long duration)
         finishCreateTableTime.add(duration, TimeUnit.NANOSECONDS);
     }
 
+    public void recordBeginCreateVectorIndexCall(long duration)
+    {
+        beginCreateVectorIndexCalls.incrementAndGet();
+        beginCreateVectorIndexTime.add(duration, TimeUnit.NANOSECONDS);
+    }
+
+    public void recordFinishCreateVectorIndexCall(long duration)
+    {
+        finishCreateVectorIndexCalls.incrementAndGet();
+        finishCreateVectorIndexTime.add(duration, TimeUnit.NANOSECONDS);
+    }
+
     public void recordGetInsertLayoutCall(long duration)
     {
         getInsertLayoutCalls.incrementAndGet();

presto-main-base/src/main/java/com/facebook/presto/metadata/StatsRecordingMetadataManager.java

@@ -29,6 +29,7 @@
 import com.facebook.presto.spi.MaterializedViewStatus;
 import com.facebook.presto.spi.MergeHandle;
 import com.facebook.presto.spi.NewTableLayout;
+import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.SystemTable;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.TableLayoutFilterCoverage;
@@ -195,6 +196,30 @@ public Optional<ConnectorOutputMetadata> finishCreateTable(Session session, Outp
         }
     }
 
+    @Override
+    public OutputTableHandle beginCreateVectorIndex(Session session, String catalogName, ConnectorTableMetadata indexMetadata, Optional<NewTableLayout> layout, SchemaTableName sourceTableName)
+    {
+        long startTime = System.nanoTime();
+        try {
+            return delegate.beginCreateVectorIndex(session, catalogName, indexMetadata, layout, sourceTableName);
+        }
+        finally {
+            stats.recordBeginCreateVectorIndexCall(System.nanoTime() - startTime);
+        }
+    }
+
+    @Override
+    public Optional<ConnectorOutputMetadata> finishCreateVectorIndex(Session session, OutputTableHandle tableHandle, Collection<Slice> fragments, Collection<ComputedStatistics> computedStatistics)
+    {
+        long startTime = System.nanoTime();
+        try {
+            return delegate.finishCreateVectorIndex(session, tableHandle, fragments, computedStatistics);
+        }
+        finally {
+            stats.recordFinishCreateVectorIndexCall(System.nanoTime() - startTime);
+        }
+    }
+
     @Override
     public Optional<NewTableLayout> getInsertLayout(Session session, TableHandle target)
     {