From ceeead8458ceb85e0ee998ad545f456271a6f2eb Mon Sep 17 00:00:00 2001 From: sunghyunkang1111 <114709653+sunghyunkang1111@users.noreply.github.com> Date: Mon, 20 May 2024 13:30:30 -0500 Subject: [PATCH] Vector search for NoSQL accounts (#1843) * Add container vector policy and indexing policy support * Add vector search capability * hide vector settings for shared throughput DB * update package-lock * fix pipeline * remove comments * Address comments * Address comments --- package-lock.json | 8 +- package.json | 2 +- src/Common/Constants.ts | 1 + .../__snapshots__/queryDocuments.test.ts.snap | 2 + src/Common/dataAccess/createCollection.ts | 5 +- src/Common/dataAccess/queryDocuments.ts | 2 + src/Contracts/DataModels.ts | 23 ++- .../CollapsibleSectionComponent.tsx | 6 +- src/Explorer/Controls/Editor/EditorReact.tsx | 11 +- .../Controls/Settings/SettingsComponent.less | 21 ++- .../Controls/Settings/SettingsComponent.tsx | 19 +++ .../ComputedPropertiesComponent.tsx | 2 +- .../ContainerVectorPolicyComponent.tsx | 30 ++++ .../IndexingPolicyComponent.tsx | 8 +- .../ComputedPropertiesComponent.test.tsx.snap | 2 +- .../IndexingPolicyComponent.test.tsx.snap | 2 +- .../Controls/Settings/SettingsUtils.tsx | 3 + .../SettingsComponent.test.tsx.snap | 1 + src/Explorer/Panes/AddCollectionPanel.tsx | 146 +++++++++++++++++- src/Explorer/Panes/PanelComponent.less | 7 + .../AddCollectionPanel.test.tsx.snap | 5 +- src/Utils/CapabilityUtils.ts | 4 + .../generatedClients/cosmos/sqlResources.ts | 2 +- .../arm/generatedClients/cosmos/types.ts | 21 +++ 24 files changed, 297 insertions(+), 36 deletions(-) create mode 100644 src/Explorer/Controls/Settings/SettingsSubComponents/ContainerVectorPolicyComponent.tsx diff --git a/package-lock.json b/package-lock.json index 8dcfe5d50..26c513d6f 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10,7 +10,7 @@ "hasInstallScript": true, "dependencies": { "@azure/arm-cosmosdb": "9.1.0", - "@azure/cosmos": "4.0.1-beta.2", + "@azure/cosmos": "4.0.1-beta.3", "@azure/cosmos-language-service": "0.0.5", "@azure/identity": "1.5.2", "@azure/ms-rest-nodeauth": "3.1.1", @@ -362,9 +362,9 @@ "integrity": "sha512-AEYxH93jGFPn/a2iVAwW87VuUIkR1FVUKB77NwMF7nBTDkDrrT/Hpt/IrCJ0QXhW27jTBDcf5ZY7w6RiqTMw2Q==" }, "node_modules/@azure/cosmos": { - "version": "4.0.1-beta.2", - "resolved": "https://registry.npmjs.org/@azure/cosmos/-/cosmos-4.0.1-beta.2.tgz", - "integrity": "sha512-iuqg/QwLQlxgRi4pnXU8JUYv+f24wkRvJ9ZZI4/sYk+DxSgkuQ194Cc2IpckpeO8z7ZpcBkVQFa82wcZVVZ8Zg==", + "version": "4.0.1-beta.3", + "resolved": "https://registry.npmjs.org/@azure/cosmos/-/cosmos-4.0.1-beta.3.tgz", + "integrity": "sha512-CpRGt+S5jnvtGUi4TmlS79YvxpbNc8/5/QHgIvvQ9D2ZFUqO0MjbMCU3lVZV2NAJT02BsbLfRAFe+FPn8nMRQw==", "dependencies": { "@azure/abort-controller": "^1.0.0", "@azure/core-auth": "^1.3.0", diff --git a/package.json b/package.json index 3a2f8b742..e1d544d1d 100644 --- a/package.json +++ b/package.json @@ -5,7 +5,7 @@ "main": "index.js", "dependencies": { "@azure/arm-cosmosdb": "9.1.0", - "@azure/cosmos": "4.0.1-beta.2", + "@azure/cosmos": "4.0.1-beta.3", "@azure/cosmos-language-service": "0.0.5", "@azure/identity": "1.5.2", "@azure/ms-rest-nodeauth": "3.1.1", diff --git a/src/Common/Constants.ts b/src/Common/Constants.ts index 0658d28c5..d56f5087a 100644 --- a/src/Common/Constants.ts +++ b/src/Common/Constants.ts @@ -88,6 +88,7 @@ export class CapabilityNames { public static readonly EnableStorageAnalytics: string = "EnableStorageAnalytics"; public static readonly EnableMongo: string = "EnableMongo"; public static readonly EnableServerless: string = "EnableServerless"; + public static readonly EnableNoSQLVectorSearch: string = "EnableNoSQLVectorSearch"; } export enum CapacityMode { diff --git a/src/Common/dataAccess/__snapshots__/queryDocuments.test.ts.snap b/src/Common/dataAccess/__snapshots__/queryDocuments.test.ts.snap index 66e740abf..ec64a72b4 100644 --- a/src/Common/dataAccess/__snapshots__/queryDocuments.test.ts.snap +++ b/src/Common/dataAccess/__snapshots__/queryDocuments.test.ts.snap @@ -2,6 +2,7 @@ exports[`getCommonQueryOptions builds the correct default options objects 1`] = ` Object { + "disableNonStreamingOrderByQuery": true, "enableScanInQuery": true, "forceQueryPlan": true, "maxDegreeOfParallelism": 0, @@ -12,6 +13,7 @@ Object { exports[`getCommonQueryOptions reads from localStorage 1`] = ` Object { + "disableNonStreamingOrderByQuery": true, "enableScanInQuery": true, "forceQueryPlan": true, "maxDegreeOfParallelism": 17, diff --git a/src/Common/dataAccess/createCollection.ts b/src/Common/dataAccess/createCollection.ts index 73ccff3f4..1d9f41844 100644 --- a/src/Common/dataAccess/createCollection.ts +++ b/src/Common/dataAccess/createCollection.ts @@ -6,13 +6,13 @@ import { Action, ActionModifiers } from "../../Shared/Telemetry/TelemetryConstan import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor"; import { userContext } from "../../UserContext"; import { getCollectionName } from "../../Utils/APITypeUtils"; +import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils"; import { createUpdateCassandraTable } from "../../Utils/arm/generatedClients/cosmos/cassandraResources"; import { createUpdateGremlinGraph } from "../../Utils/arm/generatedClients/cosmos/gremlinResources"; import { createUpdateMongoDBCollection } from "../../Utils/arm/generatedClients/cosmos/mongoDBResources"; import { createUpdateSqlContainer } from "../../Utils/arm/generatedClients/cosmos/sqlResources"; import { createUpdateTable } from "../../Utils/arm/generatedClients/cosmos/tableResources"; import * as ARMTypes from "../../Utils/arm/generatedClients/cosmos/types"; -import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils"; import { client } from "../CosmosClient"; import { handleError } from "../ErrorHandlingUtils"; import { createMongoCollectionWithProxy } from "../MongoProxyClient"; @@ -96,6 +96,9 @@ const createSqlContainer = async (params: DataModels.CreateCollectionParams): Pr if (params.uniqueKeyPolicy) { resource.uniqueKeyPolicy = params.uniqueKeyPolicy; } + if (params.vectorEmbeddingPolicy) { + resource.vectorEmbeddingPolicy = params.vectorEmbeddingPolicy; + } const rpPayload: ARMTypes.SqlDatabaseCreateUpdateParameters = { properties: { diff --git a/src/Common/dataAccess/queryDocuments.ts b/src/Common/dataAccess/queryDocuments.ts index 0b8ebd29d..223fe987d 100644 --- a/src/Common/dataAccess/queryDocuments.ts +++ b/src/Common/dataAccess/queryDocuments.ts @@ -1,4 +1,5 @@ import { FeedOptions, ItemDefinition, QueryIterator, Resource } from "@azure/cosmos"; +import { isVectorSearchEnabled } from "Utils/CapabilityUtils"; import { LocalStorageUtility, StorageKey } from "../../Shared/StorageUtility"; import { Queries } from "../Constants"; import { client } from "../CosmosClient"; @@ -26,5 +27,6 @@ export const getCommonQueryOptions = (options: FeedOptions): FeedOptions => { (storedItemPerPageSetting !== undefined && storedItemPerPageSetting) || Queries.itemsPerPage; options.maxDegreeOfParallelism = LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism); + options.disableNonStreamingOrderByQuery = !isVectorSearchEnabled(); return options; }; diff --git a/src/Contracts/DataModels.ts b/src/Contracts/DataModels.ts index e85f56964..3f7e3a7db 100644 --- a/src/Contracts/DataModels.ts +++ b/src/Contracts/DataModels.ts @@ -158,6 +158,7 @@ export interface Collection extends Resource { changeFeedPolicy?: ChangeFeedPolicy; analyticalStorageTtl?: number; geospatialConfig?: GeospatialConfig; + vectorEmbeddingPolicy?: VectorEmbeddingPolicy; schema?: ISchema; requestSchema?: () => void; computedProperties?: ComputedProperties; @@ -195,8 +196,14 @@ export interface IndexingPolicy { indexingMode: "consistent" | "lazy" | "none"; includedPaths: any; excludedPaths: any; - compositeIndexes?: any; - spatialIndexes?: any; + compositeIndexes?: any[]; + spatialIndexes?: any[]; + vectorIndexes?: VectorIndex[]; +} + +export interface VectorIndex { + path: string; + type: "flat" | "diskANN" | "quantizedFlat"; } export interface ComputedProperty { @@ -334,6 +341,18 @@ export interface CreateCollectionParams { partitionKey?: PartitionKey; uniqueKeyPolicy?: UniqueKeyPolicy; createMongoWildcardIndex?: boolean; + vectorEmbeddingPolicy?: VectorEmbeddingPolicy; +} + +export interface VectorEmbeddingPolicy { + vectorEmbeddings: VectorEmbedding[]; +} + +export interface VectorEmbedding { + dataType: "float16" | "float32" | "uint8" | "int8"; + dimensions: number; + distanceFunction: "euclidean" | "cosine" | "dotproduct"; + path: string; } export interface ReadDatabaseOfferParams { diff --git a/src/Explorer/Controls/CollapsiblePanel/CollapsibleSectionComponent.tsx b/src/Explorer/Controls/CollapsiblePanel/CollapsibleSectionComponent.tsx index d943f7cf0..207642d84 100644 --- a/src/Explorer/Controls/CollapsiblePanel/CollapsibleSectionComponent.tsx +++ b/src/Explorer/Controls/CollapsiblePanel/CollapsibleSectionComponent.tsx @@ -26,8 +26,8 @@ export class CollapsibleSectionComponent extends React.Component { @@ -75,9 +78,11 @@ export class EditorReact extends React.Component - {!this.state.showEditor && } + {!this.state.showEditor && ( + + )}
this.setRef(elt)} /> @@ -119,7 +124,7 @@ export class EditorReact extends React.Component, }); + if (this.isVectorSearchEnabled) { + tabs.push({ + tab: SettingsV2TabTypes.ContainerVectorPolicyTab, + content: , + }); + } + if (this.shouldShowIndexingPolicyEditor) { tabs.push({ tab: SettingsV2TabTypes.IndexingPolicyTab, diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/ComputedPropertiesComponent.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/ComputedPropertiesComponent.tsx index ad3b12fec..c8650b988 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/ComputedPropertiesComponent.tsx +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/ComputedPropertiesComponent.tsx @@ -121,7 +121,7 @@ export class ComputedPropertiesComponent extends React.Component<   about how to define computed properties and how to use them. -
+
); } diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/ContainerVectorPolicyComponent.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/ContainerVectorPolicyComponent.tsx new file mode 100644 index 000000000..23ccdc84c --- /dev/null +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/ContainerVectorPolicyComponent.tsx @@ -0,0 +1,30 @@ +import { Stack } from "@fluentui/react"; +import { VectorEmbeddingPolicy } from "Contracts/DataModels"; +import { EditorReact } from "Explorer/Controls/Editor/EditorReact"; +import { titleAndInputStackProps } from "Explorer/Controls/Settings/SettingsRenderUtils"; +import React from "react"; + +export interface ContainerVectorPolicyComponentProps { + vectorEmbeddingPolicy: VectorEmbeddingPolicy; +} + +export const ContainerVectorPolicyComponent: React.FC = ({ + vectorEmbeddingPolicy, +}) => { + return ( + + + + ); +}; diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/IndexingPolicyComponent.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/IndexingPolicyComponent.tsx index 4d6ca765f..bc5de93a4 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/IndexingPolicyComponent.tsx +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/IndexingPolicyComponent.tsx @@ -16,6 +16,7 @@ export interface IndexingPolicyComponentProps { logIndexingPolicySuccessMessage: () => void; indexTransformationProgress: number; refreshIndexTransformationProgress: () => Promise; + isVectorSearchEnabled?: boolean; onIndexingPolicyDirtyChange: (isIndexingPolicyDirty: boolean) => void; } @@ -119,10 +120,15 @@ export class IndexingPolicyComponent extends React.Component< indexTransformationProgress={this.props.indexTransformationProgress} refreshIndexTransformationProgress={this.props.refreshIndexTransformationProgress} /> + {this.props.isVectorSearchEnabled && ( + + Container vector policies and vector indexes are not modifiable after container creation + + )} {isDirty(this.props.indexingPolicyContent, this.props.indexingPolicyContentBaseline) && ( {unsavedEditorWarningMessage("indexPolicy")} )} -
+
); } diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/ComputedPropertiesComponent.test.tsx.snap b/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/ComputedPropertiesComponent.test.tsx.snap index 86a9bda74..2f67d7d70 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/ComputedPropertiesComponent.test.tsx.snap +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/ComputedPropertiesComponent.test.tsx.snap @@ -29,7 +29,7 @@ exports[`ComputedPropertiesComponent renders 1`] = `   about how to define computed properties and how to use them.
diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/IndexingPolicyComponent.test.tsx.snap b/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/IndexingPolicyComponent.test.tsx.snap index 1f66324f5..93516dd7e 100644 --- a/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/IndexingPolicyComponent.test.tsx.snap +++ b/src/Explorer/Controls/Settings/SettingsSubComponents/__snapshots__/IndexingPolicyComponent.test.tsx.snap @@ -12,7 +12,7 @@ exports[`IndexingPolicyComponent renders 1`] = ` refreshIndexTransformationProgress={[Function]} />
diff --git a/src/Explorer/Controls/Settings/SettingsUtils.tsx b/src/Explorer/Controls/Settings/SettingsUtils.tsx index 24ef14681..cff7d1f74 100644 --- a/src/Explorer/Controls/Settings/SettingsUtils.tsx +++ b/src/Explorer/Controls/Settings/SettingsUtils.tsx @@ -47,6 +47,7 @@ export enum SettingsV2TabTypes { IndexingPolicyTab, PartitionKeyTab, ComputedPropertiesTab, + ContainerVectorPolicyTab, } export interface IsComponentDirtyResult { @@ -152,6 +153,8 @@ export const getTabTitle = (tab: SettingsV2TabTypes): string => { return "Partition Keys (preview)"; case SettingsV2TabTypes.ComputedPropertiesTab: return "Computed Properties"; + case SettingsV2TabTypes.ContainerVectorPolicyTab: + return "Container Vector Policy (preview)"; default: throw new Error(`Unknown tab ${tab}`); } diff --git a/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap b/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap index b536bed8a..561368bdd 100644 --- a/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap +++ b/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap @@ -196,6 +196,7 @@ exports[`SettingsComponent renders 1`] = ` "indexingMode": "consistent", } } + isVectorSearchEnabled={false} logIndexingPolicySuccessMessage={[Function]} onIndexingPolicyContentChange={[Function]} onIndexingPolicyDirtyChange={[Function]} diff --git a/src/Explorer/Panes/AddCollectionPanel.tsx b/src/Explorer/Panes/AddCollectionPanel.tsx index 8c1cfef5b..7bb0f553d 100644 --- a/src/Explorer/Panes/AddCollectionPanel.tsx +++ b/src/Explorer/Panes/AddCollectionPanel.tsx @@ -21,6 +21,7 @@ import { getErrorMessage, getErrorStack } from "Common/ErrorHandlingUtils"; import { configContext, Platform } from "ConfigContext"; import * as DataModels from "Contracts/DataModels"; import { SubscriptionType } from "Contracts/SubscriptionType"; +import { EditorReact } from "Explorer/Controls/Editor/EditorReact"; import { useSidePanel } from "hooks/useSidePanel"; import { useTeachingBubble } from "hooks/useTeachingBubble"; import React from "react"; @@ -29,7 +30,7 @@ import { Action } from "Shared/Telemetry/TelemetryConstants"; import * as TelemetryProcessor from "Shared/Telemetry/TelemetryProcessor"; import { userContext } from "UserContext"; import { getCollectionName } from "Utils/APITypeUtils"; -import { isCapabilityEnabled, isServerlessAccount } from "Utils/CapabilityUtils"; +import { isCapabilityEnabled, isServerlessAccount, isVectorSearchEnabled } from "Utils/CapabilityUtils"; import { getUpsellMessage } from "Utils/PricingUtils"; import { CollapsibleSectionComponent } from "../Controls/CollapsiblePanel/CollapsibleSectionComponent"; import { ThroughputInput } from "../Controls/ThroughputInput/ThroughputInput"; @@ -81,6 +82,26 @@ export const AllPropertiesIndexed: DataModels.IndexingPolicy = { excludedPaths: [], }; +const DefaultDatabaseVectorIndex: DataModels.IndexingPolicy = { + indexingMode: "consistent", + automatic: true, + includedPaths: [ + { + path: "/*", + }, + ], + excludedPaths: [ + { + path: '/"_etag"/?', + }, + ], + vectorIndexes: [], +}; + +export const DefaultVectorEmbeddingPolicy: DataModels.VectorEmbeddingPolicy = { + vectorEmbeddings: [], +}; + export interface AddCollectionPanelState { createNewDatabase: boolean; newDatabaseId: string; @@ -101,6 +122,8 @@ export interface AddCollectionPanelState { isExecuting: boolean; isThroughputCapExceeded: boolean; teachingBubbleStep: number; + vectorIndexingPolicy: string; + vectorEmbeddingPolicy: string; } export class AddCollectionPanel extends React.Component { @@ -136,6 +159,8 @@ export class AddCollectionPanel extends React.Component +
{this.state.errorMessage && ( )} - + {this.shouldShowVectorSearchParameters() && ( + + { + this.scrollToSection("collapsibleVectorPolicySectionContent"); + }} + > + + + Learn more + + this.setVectorIndexingPolicy(newIndexingPolicy)} + /> + + + { + this.scrollToSection("collapsibleVectorPolicySectionContent"); + }} + > + + + Learn more + + + this.setVectorEmbeddingPolicy(newVectorEmbeddingPolicy) + } + /> + + + + )} {userContext.apiType !== "Tables" && ( { TelemetryProcessor.traceOpen(Action.ExpandAddCollectionPaneAdvancedSection); - this.scrollToAdvancedSection(); + this.scrollToSection("collapsibleAdvancedSectionContent"); }} > - + {isCapabilityEnabled("EnableMongo") && !isCapabilityEnabled("EnableMongo16MBDocumentSupport") && ( @@ -1070,6 +1160,18 @@ export class AddCollectionPanel extends React.Component
-`; \ No newline at end of file +`; diff --git a/src/Utils/CapabilityUtils.ts b/src/Utils/CapabilityUtils.ts index 216eff13d..8b6976666 100644 --- a/src/Utils/CapabilityUtils.ts +++ b/src/Utils/CapabilityUtils.ts @@ -16,3 +16,7 @@ export const isServerlessAccount = (): boolean => { isCapabilityEnabled(Constants.CapabilityNames.EnableServerless) ); }; + +export const isVectorSearchEnabled = (): boolean => { + return userContext.apiType === "SQL" && isCapabilityEnabled(Constants.CapabilityNames.EnableNoSQLVectorSearch); +}; diff --git a/src/Utils/arm/generatedClients/cosmos/sqlResources.ts b/src/Utils/arm/generatedClients/cosmos/sqlResources.ts index f85ddb636..049e265e9 100644 --- a/src/Utils/arm/generatedClients/cosmos/sqlResources.ts +++ b/src/Utils/arm/generatedClients/cosmos/sqlResources.ts @@ -9,7 +9,7 @@ import { configContext } from "../../../../ConfigContext"; import { armRequest } from "../../request"; import * as Types from "./types"; -const apiVersion = "2024-02-15-preview"; +const apiVersion = "2024-05-15-preview"; /* Lists the SQL databases under an existing Azure Cosmos DB database account. */ export async function listSqlDatabases( diff --git a/src/Utils/arm/generatedClients/cosmos/types.ts b/src/Utils/arm/generatedClients/cosmos/types.ts index 4f69d223c..5272f215b 100644 --- a/src/Utils/arm/generatedClients/cosmos/types.ts +++ b/src/Utils/arm/generatedClients/cosmos/types.ts @@ -1235,6 +1235,9 @@ export interface SqlDatabaseResource { export interface SqlContainerResource { /* Name of the Cosmos DB SQL container */ id: string; + + vectorEmbeddingPolicy?: VectorEmbeddingPolicy; + /* The configuration of the indexing policy. By default, the indexing is automatic for all document paths within the container */ indexingPolicy?: IndexingPolicy; @@ -1267,6 +1270,17 @@ export interface SqlContainerResource { computedProperties?: ComputedProperty[]; } +export interface VectorEmbeddingPolicy { + vectorEmbeddings: VectorEmbedding[]; +} + +export interface VectorEmbedding { + path?: string; + dataType?: string; + dimensions?: number; + distanceFunction?: string; +} + /* Cosmos DB indexing policy */ export interface IndexingPolicy { /* Indicates if the indexing policy is automatic */ @@ -1285,6 +1299,13 @@ export interface IndexingPolicy { /* List of spatial specifics */ spatialIndexes?: SpatialSpec[]; + + vectorIndexes?: VectorIndex[]; +} + +export interface VectorIndex { + path?: string; + type?: string; } /* undocumented */