Merge branch 'main' into expr-union-type-impl

Author: chardoncs

.github/workflows/_doc_release.yml

@@ -0,0 +1,31 @@
+name: Release Docs
+
+on:
+  workflow_call:
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    environment: docs-release
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v4
+        with:
+          node-version: 18
+          cache: yarn
+          cache-dependency-path: docs/yarn.lock
+      - uses: webfactory/[email protected]
+        with:
+          ssh-private-key: ${{ secrets.GH_PAGES_DEPLOY }}
+      - name: Deploy to GitHub Pages
+        env:
+          USE_SSH: true
+        run: |
+          export COCOINDEX_DOCS_POSTHOG_API_KEY=${{ vars.COCOINDEX_DOCS_POSTHOG_API_KEY }}
+          export COCOINDEX_DOCS_MIXPANEL_API_KEY=${{ vars.COCOINDEX_DOCS_MIXPANEL_API_KEY }}
+          export COCOINDEX_DOCS_ALGOLIA_APP_ID=${{ vars.COCOINDEX_DOCS_ALGOLIA_APP_ID }}
+          export COCOINDEX_DOCS_ALGOLIA_API_KEY=${{ vars.COCOINDEX_DOCS_ALGOLIA_API_KEY }}
+          git config --global user.email "${{ vars.COCOINDEX_DOCS_DEPLOY_USER_EMAIL }}"
+          git config --global user.name "${{ vars.COCOINDEX_DOCS_DEPLOY_USER_NAME }}"
+          yarn --cwd docs install --frozen-lockfile
+          yarn --cwd docs deploy

.github/workflows/docs.yml

@@ -6,19 +6,14 @@ on:
     paths:
       - docs/**
       - ".github/workflows/docs.yml"
-  push:
-    branches: [main]
-    paths:
-      - docs/**
-      - ".github/workflows/docs.yml"
   workflow_dispatch:
 
 permissions:
   contents: write
 
 jobs:
   test-deploy:
-    if: github.event_name != 'push'
+    if: github.event_name == 'pull_request'
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -32,43 +27,8 @@ jobs:
       - name: Test build website
         run: yarn --cwd docs build
 
-  deploy-precheck:
-    if: ${{ github.event_name != 'pull_request' }}
-    runs-on: ubuntu-latest
-    environment: docs-release
-    outputs:
-      gh-deploy-key: ${{ steps.gh-deploy-key.outputs.defined }}
-    steps:
-        - id: gh-deploy-key
-          env:
-              GH_PAGES_DEPLOY: ${{ secrets.GH_PAGES_DEPLOY }}
-          if: "${{ env.GH_PAGES_DEPLOY != '' }}"
-          run: echo "defined=true" >> $GITHUB_OUTPUT
-
   deploy:
-    needs: [deploy-precheck]
-    if: ${{ needs.deploy-precheck.outputs.gh-deploy-key == 'true' }}
-    runs-on: ubuntu-latest
-    environment: docs-release
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-node@v4
-        with:
-          node-version: 18
-          cache: yarn
-          cache-dependency-path: docs/yarn.lock
-      - uses: webfactory/[email protected]
-        with:
-          ssh-private-key: ${{ secrets.GH_PAGES_DEPLOY }}
-      - name: Deploy to GitHub Pages
-        env:
-          USE_SSH: true
-        run: |
-          export COCOINDEX_DOCS_POSTHOG_API_KEY=${{ vars.COCOINDEX_DOCS_POSTHOG_API_KEY }}
-          export COCOINDEX_DOCS_MIXPANEL_API_KEY=${{ vars.COCOINDEX_DOCS_MIXPANEL_API_KEY }}
-          export COCOINDEX_DOCS_ALGOLIA_APP_ID=${{ vars.COCOINDEX_DOCS_ALGOLIA_APP_ID }}
-          export COCOINDEX_DOCS_ALGOLIA_API_KEY=${{ vars.COCOINDEX_DOCS_ALGOLIA_API_KEY }}
-          git config --global user.email "${{ vars.COCOINDEX_DOCS_DEPLOY_USER_EMAIL }}"
-          git config --global user.name "${{ vars.COCOINDEX_DOCS_DEPLOY_USER_NAME }}"
-          yarn --cwd docs install --frozen-lockfile
-          yarn --cwd docs deploy
+    name: Release Docs
+    if: ${{ github.event_name == 'workflow_dispatch' }}
+    uses: ./.github/workflows/_doc_release.yml
+    secrets: inherit

.github/workflows/release.yml

@@ -107,3 +107,9 @@ jobs:
         with:
           command: upload
           args: --non-interactive --skip-existing wheels-*/*
+    
+  release-docs:
+    name: Release Docs
+    needs: [release]
+    uses: ./.github/workflows/_doc_release.yml
+    secrets: inherit

.gitignore

@@ -9,7 +9,7 @@ __pycache__/
 *.so
 
 # Distribution / packaging
-.venv/
+.venv*/
 dist/
 
 .DS_Store

Cargo.lock

@@ -4,6 +4,7 @@ name = "cocoindex"
 # Will be overridden for specific release versions.
 version = "999.0.0"
 edition = "2024"
+rust-version = "1.86"
 
 [profile.release]
 codegen-units = 1
@@ -50,6 +51,7 @@ tower-http = { version = "0.6.2", features = ["cors", "trace"] }
 indexmap = { version = "2.8.0", features = ["serde"] }
 blake2 = "0.10.6"
 pgvector = { version = "0.4.0", features = ["sqlx"] }
+phf = { version = "0.11.3", features = ["macros"] }
 indenter = "0.3.3"
 itertools = "0.14.0"
 derivative = "2.2.0"

Cargo.toml

@@ -36,16 +36,17 @@ This is the list of all basic types supported by CocoIndex:
 | LocalDatetime | Date and time without timezone | `cocoindex.LocalDateTime` | `datetime.datetime` |
 | OffsetDatetime | Date and time with a timezone offset | `cocoindex.OffsetDateTime` | `datetime.datetime` |
 | TimeDelta | A duration of time | `datetime.timedelta` | `datetime.timedelta` |
-| Vector[*T*, *Dim*?] | *T* must be basic type. *Dim* is a positive integer and optional. |`cocoindex.Vector[T]` or `cocoindex.Vector[T, Dim]` | `list[T]` | 
 | Json | | `cocoindex.Json` | Any data convertible to JSON by `json` package | 
+| Vector[*T*, *Dim*?] | *T* can be a basic type or a numeric type. *Dim* is a positive integer and optional. | `cocoindex.Vector[T]` or `cocoindex.Vector[T, Dim]` | `numpy.typing.NDArray[T]` or `list[T]` |
 
 Values of all data types can be represented by values in Python's native types (as described under the Native Python Type column).
 However, the underlying execution engine and some storage system (like Postgres) has finer distinctions for some types, specifically:
 
 *   *Float32* and *Float64* for `float`, with different precision.
 *   *LocalDateTime* and *OffsetDateTime* for `datetime.datetime`, with different timezone awareness.
-*   *Vector* has optional dimension information.
 *   *Range* and *Json* provide a clear tag for the type, to clearly distinguish the type in CocoIndex.
+*   *Vector* holds elements of type *T*. If *T* is numeric (e.g., `np.float32` or `np.float64`), it's represented as `NDArray[T]`; otherwise, as `list[T]`.
+*   *Vector* also has optional dimension information.
 
 The native Python type is always more permissive and can represent a superset of possible values.
 *   Only when you annotate the return type of a custom function, you should use the specific type,

docs/docs/core/data_types.mdx

@@ -154,11 +154,11 @@ The goal of transforming your data is usually to query against it.
 Once you already have your index built, you can directly access the transformed data in the target database.
 CocoIndex also provides utilities for you to do this more seamlessly.
 
-In this example, we'll use the [`psycopg` library](https://www.psycopg.org/) to connect to the database and run queries.
-Please make sure it's installed:
+In this example, we'll use the [`psycopg` library](https://www.psycopg.org/) along with pgvector to connect to the database and run queries on vector data.
+Please make sure the required packages are installed:
 
 ```bash
-pip install psycopg[binary,pool]
+pip install numpy "psycopg[binary,pool]" pgvector
 ```
 
 ### Step 4.1: Extract common transformations
@@ -169,8 +169,11 @@ i.e. they should use exactly the same embedding model and parameters.
 Let's extract that into a function:
 
 ```python title="quickstart.py"
+from numpy.typing import NDArray
+import numpy as np
+
 @cocoindex.transform_flow()
-def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
+def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[NDArray[np.float32]]:
     return text.transform(
         cocoindex.functions.SentenceTransformerEmbed(
             model="sentence-transformers/all-MiniLM-L6-v2"))
@@ -207,6 +210,7 @@ Now we can create a function to query the index upon a given input query:
 
 ```python title="quickstart.py"
 from psycopg_pool import ConnectionPool
+from pgvector.psycopg import register_vector
 
 def search(pool: ConnectionPool, query: str, top_k: int = 5):
     # Get the table name, for the export target in the text_embedding_flow above.
@@ -215,9 +219,10 @@ def search(pool: ConnectionPool, query: str, top_k: int = 5):
     query_vector = text_to_embedding.eval(query)
     # Run the query and get the results.
     with pool.connection() as conn:
+        register_vector(conn)
         with conn.cursor() as cur:
             cur.execute(f"""
-                SELECT filename, text, embedding <=> %s::vector AS distance
+                SELECT filename, text, embedding <=> %s AS distance
                 FROM {table_name} ORDER BY distance LIMIT %s
             """, (query_vector, top_k))
             return [
@@ -236,7 +241,7 @@ There're two CocoIndex-specific logic:
 
 2.  Evaluate the transform flow defined above with the input query, to get the embedding.
     It's done by the `eval()` method of the transform flow `text_to_embedding`.
-    The return type of this method is `list[float]` as declared in the `text_to_embedding()` function (`cocoindex.DataSlice[list[float]]`).
+    The return type of this method is `NDArray[np.float32]` as declared in the `text_to_embedding()` function (`cocoindex.DataSlice[NDArray[np.float32]]`).
 
 ### Step 4.3: Add the main script logic
 

docs/docs/getting_started/quickstart.md

@@ -41,7 +41,7 @@ The [quickstart](getting_started/quickstart#step-41-extract-common-transformatio
 
 ```python
 @cocoindex.transform_flow()
-def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
+def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[NDArray[np.float32]]:
     return text.transform(
         cocoindex.functions.SentenceTransformerEmbed(
             model="sentence-transformers/all-MiniLM-L6-v2"))
@@ -61,7 +61,7 @@ with doc["chunks"].row() as chunk:
     chunk["embedding"] = chunk["text"].call(text_to_embedding)
 ```
 
-Any time, you can call the `eval()` method with specific string, which will return a `list[float]`:
+Any time, you can call the `eval()` method with specific string, which will return a `NDArray[np.float32]`:
 
 ```python
 print(text_to_embedding.eval("Hello, world!"))
@@ -93,7 +93,7 @@ For example:
 
 ```python
 table_name = cocoindex.utils.get_target_storage_default_name(text_embedding_flow, "doc_embeddings")
-query = f"SELECT filename, text FROM {table_name} ORDER BY embedding <=> %s::vector DESC LIMIT 5"
+query = f"SELECT filename, text FROM {table_name} ORDER BY embedding <=> %s DESC LIMIT 5"
 ...
 ```
 

docs/docs/query.mdx

@@ -45,7 +45,7 @@
    },
    "outputs": [],
    "source": [
-    "%pip install cocoindex python-dotenv psycopg[binary,pool]"
+    "%pip install cocoindex numpy python-dotenv psycopg[binary,pool] pgvector"
    ]
   },
   {
@@ -164,7 +164,10 @@
     "from dotenv import load_dotenv\n",
     "import os\n",
     "from psycopg_pool import ConnectionPool\n",
-    "import cocoindex\n"
+    "from pgvector.psycopg import register_vector\n",
+    "import cocoindex\n",
+    "from numpy.typing import NDArray\n",
+    "import numpy as np\n"
    ]
   },
   {
@@ -187,7 +190,7 @@
     "%%writefile -a main.py\n",
     "\n",
     "@cocoindex.transform_flow()\n",
-    "def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:\n",
+    "def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[NDArray[np.float32]]:\n",
     "    \"\"\"\n",
     "    Embed the text using a SentenceTransformer model.\n",
     "    This is shared logic between indexing and querying.\n",
@@ -274,9 +277,10 @@
     "    query_vector = text_to_embedding.eval(query)\n",
     "    # Run the query and get the results.\n",
     "    with pool.connection() as conn:\n",
+    "        register_vector(conn)\n",
     "        with conn.cursor() as cur:\n",
     "            cur.execute(f\"\"\"\n",
-    "                SELECT filename, text, embedding <=> %s::vector AS distance\n",
+    "                SELECT filename, text, embedding <=> %s AS distance\n",
     "                FROM {table_name} ORDER BY distance LIMIT %s\n",
     "            \"\"\", (query_vector, top_k))\n",
     "            return [\n",