Bump version: 0.8.2-beta.0 → 0.8.2

This also exposed the max_batch_length configuration option in
python/node (it was needed to verify if we are actually in v2 mode or
not)

.bumpversion.cfg

@@ -1,22 +0,0 @@
-[bumpversion]
-current_version = 0.4.20
-commit = True
-message = Bump version: {current_version} → {new_version}
-tag = True
-tag_name = v{new_version}
-
-[bumpversion:file:node/package.json]
-
-[bumpversion:file:nodejs/package.json]
-
-[bumpversion:file:nodejs/npm/darwin-x64/package.json]
-
-[bumpversion:file:nodejs/npm/darwin-arm64/package.json]
-
-[bumpversion:file:nodejs/npm/linux-x64-gnu/package.json]
-
-[bumpversion:file:nodejs/npm/linux-arm64-gnu/package.json]
-
-[bumpversion:file:rust/ffi/node/Cargo.toml]
-
-[bumpversion:file:rust/lancedb/Cargo.toml]

.bumpversion.toml

@@ -0,0 +1,57 @@
+[tool.bumpversion]
+current_version = "0.5.1"
+parse = """(?x)
+    (?P<major>0|[1-9]\\d*)\\.
+    (?P<minor>0|[1-9]\\d*)\\.
+    (?P<patch>0|[1-9]\\d*)
+    (?:-(?P<pre_l>[a-zA-Z-]+)\\.(?P<pre_n>0|[1-9]\\d*))?
+"""
+serialize = [
+    "{major}.{minor}.{patch}-{pre_l}.{pre_n}",
+    "{major}.{minor}.{patch}",
+]
+search = "{current_version}"
+replace = "{new_version}"
+regex = false
+ignore_missing_version = false
+ignore_missing_files = false
+tag = true
+sign_tags = false
+tag_name = "v{new_version}"
+tag_message = "Bump version: {current_version} → {new_version}"
+allow_dirty = true
+commit = true
+message = "Bump version: {current_version} → {new_version}"
+commit_args = ""
+
+[tool.bumpversion.parts.pre_l]
+values = ["beta", "final"]
+optional_value = "final"
+
+[[tool.bumpversion.files]]
+filename = "node/package.json"
+search = "\"version\": \"{current_version}\","
+replace = "\"version\": \"{new_version}\","
+
+[[tool.bumpversion.files]]
+filename = "nodejs/package.json"
+search = "\"version\": \"{current_version}\","
+replace = "\"version\": \"{new_version}\","
+
+# nodejs binary packages
+[[tool.bumpversion.files]]
+glob = "nodejs/npm/*/package.json"
+search = "\"version\": \"{current_version}\","
+replace = "\"version\": \"{new_version}\","
+
+# Cargo files
+# ------------
+[[tool.bumpversion.files]]
+filename = "rust/ffi/node/Cargo.toml"
+search = "\nversion = \"{current_version}\""
+replace = "\nversion = \"{new_version}\""
+
+[[tool.bumpversion.files]]
+filename = "rust/lancedb/Cargo.toml"
+search = "\nversion = \"{current_version}\""
+replace = "\nversion = \"{new_version}\""

.github/release.yml

@@ -1,25 +0,0 @@
-# TODO: create separate templates for Python and other releases.
-changelog:
-  exclude:
-    labels:
-      - ci
-      - chore
-  categories:
-    - title: Breaking Changes 🛠
-      labels:
-        - breaking-change
-    - title: New Features 🎉
-      labels:
-        - enhancement
-    - title: Bug Fixes 🐛
-      labels:
-        - bug
-    - title: Documentation 📚
-      labels:
-        - documentation
-    - title: Performance Improvements 🚀
-      labels:
-        - performance
-    - title: Other Changes
-      labels:
-        - "*"

.github/release_notes.json

@@ -0,0 +1,41 @@
+{
+    "ignore_labels": ["chore"],
+    "pr_template": "- ${{TITLE}} by @${{AUTHOR}} in ${{URL}}",
+    "categories": [
+        {
+            "title": "## 🏆 Highlights",
+            "labels": ["highlight"]
+        },
+        {
+            "title": "## 🛠 Breaking Changes",
+            "labels": ["breaking-change"]
+        },
+        {
+            "title": "## ⚠️ Deprecations ",
+            "labels": ["deprecation"]
+        },
+        {
+            "title": "## 🎉 New Features",
+            "labels": ["enhancement"]
+        },
+        {
+            "title": "## 🐛 Bug Fixes",
+            "labels": ["bug"]
+        },
+        {
+            "title": "## 📚 Documentation",
+            "labels": ["documentation"]
+        },
+        {
+            "title": "## 🚀 Performance Improvements",
+            "labels": ["performance"]
+        },
+        {
+            "title": "## Other Changes"
+        },
+        {
+            "title": "## 🔧 Build and CI",
+            "labels": ["ci"]
+        }
+    ]
+}

.github/workflows/cargo-publish.yml

@@ -1,8 +1,12 @@
 name: Cargo Publish
 
 on:
-  release:
-    types: [ published ]
+  push:
+    tags-ignore:
+      # We don't publish pre-releases for Rust. Crates.io is just a source
+      # distribution, so we don't need to publish pre-releases.
+      - 'v*-beta*'
+      - '*-v*' # for example, python-vX.Y.Z
 
 env:
   # This env var is used by Swatinem/rust-cache@v2 for the cache

.github/workflows/java.yml

@@ -0,0 +1,85 @@
+name: Build and Run Java JNI Tests
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    paths:
+      - java/**
+      - rust/**
+      - .github/workflows/java.yml
+env:
+  # This env var is used by Swatinem/rust-cache@v2 for the cache
+  # key, so we set it to make sure it is always consistent.
+  CARGO_TERM_COLOR: always
+  # Disable full debug symbol generation to speed up CI build and keep memory down
+  # "1" means line tables only, which is useful for panic tracebacks.
+  RUSTFLAGS: "-C debuginfo=1"
+  RUST_BACKTRACE: "1"
+  # according to: https://matklad.github.io/2021/09/04/fast-rust-builds.html
+  # CI builds are faster with incremental disabled.
+  CARGO_INCREMENTAL: "0"
+  CARGO_BUILD_JOBS: "1"
+jobs:
+  linux-build:
+    runs-on: ubuntu-22.04
+    name: ubuntu-22.04 + Java 11 & 17
+    defaults:
+      run:
+        working-directory: ./java
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - uses: Swatinem/rust-cache@v2
+        with:
+          workspaces: java/core/lancedb-jni
+      - name: Run cargo fmt
+        run: cargo fmt --check
+        working-directory: ./java/core/lancedb-jni
+      - name: Install dependencies
+        run: |
+          sudo apt update
+          sudo apt install -y protobuf-compiler libssl-dev
+      - name: Install Java 17
+        uses: actions/setup-java@v4
+        with:
+          distribution: temurin
+          java-version: 17
+          cache: "maven"
+      - run: echo "JAVA_17=$JAVA_HOME" >> $GITHUB_ENV
+      - name: Install Java 11
+        uses: actions/setup-java@v4
+        with:
+          distribution: temurin
+          java-version: 11
+          cache: "maven"
+      - name: Java Style Check
+        run: mvn checkstyle:check
+      # Disable because of issues in lancedb rust core code
+      # - name: Rust Clippy
+      #   working-directory: java/core/lancedb-jni
+      #   run: cargo clippy --all-targets -- -D warnings
+      - name: Running tests with Java 11
+        run: mvn clean test
+      - name: Running tests with Java 17
+        run: |
+          export JAVA_TOOL_OPTIONS="$JAVA_TOOL_OPTIONS \
+          -XX:+IgnoreUnrecognizedVMOptions \
+          --add-opens=java.base/java.lang=ALL-UNNAMED \
+          --add-opens=java.base/java.lang.invoke=ALL-UNNAMED \
+          --add-opens=java.base/java.lang.reflect=ALL-UNNAMED \
+          --add-opens=java.base/java.io=ALL-UNNAMED \
+          --add-opens=java.base/java.net=ALL-UNNAMED \
+          --add-opens=java.base/java.nio=ALL-UNNAMED \
+          --add-opens=java.base/java.util=ALL-UNNAMED \
+          --add-opens=java.base/java.util.concurrent=ALL-UNNAMED \
+          --add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED \
+          --add-opens=java.base/jdk.internal.ref=ALL-UNNAMED \
+          --add-opens=java.base/sun.nio.ch=ALL-UNNAMED \
+          --add-opens=java.base/sun.nio.cs=ALL-UNNAMED \
+          --add-opens=java.base/sun.security.action=ALL-UNNAMED \
+          --add-opens=java.base/sun.util.calendar=ALL-UNNAMED \
+          --add-opens=java.security.jgss/sun.security.krb5=ALL-UNNAMED \
+          -Djdk.reflect.useDirectMethodHandle=false \
+          -Dio.netty.tryReflectionSetAccessible=true"
+          JAVA_HOME=$JAVA_17 mvn clean test

.github/workflows/make-release-commit.yml

@@ -1,37 +1,62 @@
 name: Create release commit
 
+# This workflow increments versions, tags the version, and pushes it.
+# When a tag is pushed, another workflow is triggered that creates a GH release
+# and uploads the binaries. This workflow is only for creating the tag.
+
+# This script will enforce that a minor version is incremented if there are any
+# breaking changes since the last minor increment. However, it isn't able to
+# differentiate between breaking changes in Node versus Python. If you wish to
+# bypass this check, you can manually increment the version and push the tag.
 on:
   workflow_dispatch:
     inputs:
       dry_run:
         description: 'Dry run (create the local commit/tags but do not push it)'
         required: true
-        default: "false"
-        type: choice
-        options:
-          - "true"
-          - "false"
-      part:
+        default: false
+        type: boolean
+      type:
         description: 'What kind of release is this?'
         required: true
-        default: 'patch'
+        default: 'preview'
         type: choice
         options:
-          - patch
-          - minor
-          - major
+          - preview
+          - stable
+      python:
+        description: 'Make a Python release'
+        required: true
+        default: true
+        type: boolean
+      other:
+        description: 'Make a Node/Rust release'
+        required: true
+        default: true
+        type: boolean
+      bump-minor:
+        description: 'Bump minor version'
+        required: true
+        default: false
+        type: boolean
 
 jobs:
-  bump-version:
+  make-release:
+    # Creates tag and GH release. The GH release will trigger the build and release jobs.
     runs-on: ubuntu-latest
+    permissions:
+      contents: write
     steps:
-      - name: Check out main
-        uses: actions/checkout@v4
+      - name: Output Inputs
+        run: echo "${{ toJSON(github.event.inputs) }}"
+      - uses: actions/checkout@v4
         with:
-          ref: main
-          persist-credentials: false
           fetch-depth: 0
           lfs: true
+          # It's important we use our token here, as the default token will NOT
+          # trigger any workflows watching for new tags. See:
+          # https://docs.github.com/en/actions/using-workflows/triggering-a-workflow#triggering-a-workflow-from-a-workflow
+          token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
       - name: Set git configs for bumpversion
         shell: bash
         run: |
@@ -41,19 +66,34 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: "3.11"
-      - name: Bump version, create tag and commit
+      - name: Bump Python version
+        if: ${{ inputs.python }}
+        working-directory: python
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: |
-          pip install bump2version
-          bumpversion --verbose ${{ inputs.part }}
-      - name: Push new version and tag
-        if: ${{ inputs.dry_run }} == "false"
+          # Need to get the commit before bumping the version, so we can
+          # determine if there are breaking changes in the next step as well.
+          echo "COMMIT_BEFORE_BUMP=$(git rev-parse HEAD)" >> $GITHUB_ENV
+
+          pip install bump-my-version PyGithub packaging
+          bash ../ci/bump_version.sh ${{ inputs.type }} ${{ inputs.bump-minor }} python-v
+      - name: Bump Node/Rust version
+        if: ${{ inputs.other }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          pip install bump-my-version PyGithub packaging
+          bash ci/bump_version.sh ${{ inputs.type }} ${{ inputs.bump-minor }} v $COMMIT_BEFORE_BUMP
+      - name: Push new version tag
+        if: ${{ !inputs.dry_run }}
         uses: ad-m/github-push-action@master
         with:
+          # Need to use PAT here too to trigger next workflow. See comment above.
           github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
-          branch: main
+          branch: ${{ github.ref }}
           tags: true
       - uses: ./.github/workflows/update_package_lock
-        if: ${{ inputs.dry_run }} == "false"
+        if: ${{ !inputs.dry_run && inputs.other }}
         with:
-          github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
-
+          github_token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/npm-publish.yml

@@ -1,8 +1,9 @@
 name: NPM Publish
 
 on:
-  release:
-    types: [published]
+  push:
+    tags:
+      - 'v*'
 
 jobs:
   node:
@@ -274,9 +275,15 @@ jobs:
         env:
           NODE_AUTH_TOKEN: ${{ secrets.LANCEDB_NPM_REGISTRY_TOKEN }}
         run: |
+          # Tag beta as "preview" instead of default "latest". See lancedb 
+          # npm publish step for more info.
+          if [[ $GITHUB_REF =~ refs/tags/v(.*)-beta.* ]]; then
+            PUBLISH_ARGS="--tag preview"
+          fi
+
           mv */*.tgz .
           for filename in *.tgz; do
-            npm publish $filename
+            npm publish $PUBLISH_ARGS $filename
           done
 
   release-nodejs:
@@ -316,11 +323,23 @@ jobs:
       - name: Publish to NPM
         env:
           NODE_AUTH_TOKEN: ${{ secrets.LANCEDB_NPM_REGISTRY_TOKEN }}
-        run: npm publish --access public
+        # By default, things are published to the latest tag. This is what is 
+        # installed by default if the user does not specify a version. This is
+        # good for stable releases, but for pre-releases, we want to publish to
+        # the "preview" tag so they can install with `npm install lancedb@preview`.
+        # See: https://medium.com/@mbostock/prereleases-and-npm-e778fc5e2420
+        run: |
+          if [[ $GITHUB_REF =~ refs/tags/v(.*)-beta.* ]]; then
+            npm publish --access public --tag preview
+          else
+            npm publish --access public
+          fi
 
   update-package-lock:
     needs: [release]
     runs-on: ubuntu-latest
+    permissions:
+      contents: write
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -331,11 +350,13 @@ jobs:
           lfs: true
       - uses: ./.github/workflows/update_package_lock
         with:
-          github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
+          github_token: ${{ secrets.GITHUB_TOKEN }}
 
   update-package-lock-nodejs:
     needs: [release-nodejs]
     runs-on: ubuntu-latest
+    permissions:
+      contents: write
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -346,4 +367,70 @@ jobs:
           lfs: true
       - uses: ./.github/workflows/update_package_lock_nodejs
         with:
-          github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+  
+  gh-release:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Extract version
+        id: extract_version
+        env:
+          GITHUB_REF: ${{ github.ref }}
+        run: |
+          set -e
+          echo "Extracting tag and version from $GITHUB_REF"
+          if [[ $GITHUB_REF =~ refs/tags/v(.*) ]]; then
+            VERSION=${BASH_REMATCH[1]}
+            TAG=v$VERSION
+            echo "tag=$TAG" >> $GITHUB_OUTPUT
+            echo "version=$VERSION" >> $GITHUB_OUTPUT
+          else
+            echo "Failed to extract version from $GITHUB_REF"
+            exit 1
+          fi
+          echo "Extracted version $VERSION from $GITHUB_REF"
+          if [[ $VERSION =~ beta ]]; then
+            echo "This is a beta release"
+
+            # Get last release (that is not this one)
+            FROM_TAG=$(git tag --sort='version:refname' \
+              | grep ^v \
+              | grep -vF "$TAG" \
+              | python ci/semver_sort.py v \
+              | tail -n 1)
+          else
+            echo "This is a stable release"
+            # Get last stable tag (ignore betas)
+            FROM_TAG=$(git tag --sort='version:refname' \
+              | grep ^v \
+              | grep -vF "$TAG" \
+              | grep -v beta \
+              | python ci/semver_sort.py v \
+              | tail -n 1)
+          fi
+          echo "Found from tag $FROM_TAG"
+          echo "from_tag=$FROM_TAG" >> $GITHUB_OUTPUT
+      - name: Create Release Notes
+        id: release_notes
+        uses: mikepenz/release-changelog-builder-action@v4
+        with:
+          configuration: .github/release_notes.json
+          toTag: ${{ steps.extract_version.outputs.tag }}
+          fromTag: ${{ steps.extract_version.outputs.from_tag }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      - name: Create GH release
+        uses: softprops/action-gh-release@v2
+        with:
+          prerelease: ${{ contains('beta', github.ref) }}
+          tag_name: ${{ steps.extract_version.outputs.tag }}
+          token: ${{ secrets.GITHUB_TOKEN }}
+          generate_release_notes: false
+          name: Node/Rust LanceDB v${{ steps.extract_version.outputs.version }}
+          body: ${{ steps.release_notes.outputs.changelog }}

.github/workflows/pypi-publish.yml

@@ -1,18 +1,16 @@
 name: PyPI Publish
 
 on:
-  release:
-    types: [published]
+  push:
+    tags:
+      - 'python-v*'
 
 jobs:
   linux:
-    # Only runs on tags that matches the python-make-release action
-    if: startsWith(github.ref, 'refs/tags/python-v')
     name: Python ${{ matrix.config.platform }} manylinux${{ matrix.config.manylinux }}
     timeout-minutes: 60
     strategy:
       matrix:
-        python-minor-version: ["8"]
         config:
           - platform: x86_64
             manylinux: "2_17"
@@ -34,25 +32,22 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: 3.${{ matrix.python-minor-version }}
+          python-version: 3.8
       - uses: ./.github/workflows/build_linux_wheel
         with:
-          python-minor-version: ${{ matrix.python-minor-version }}
+          python-minor-version: 8
           args: "--release --strip ${{ matrix.config.extra_args }}"
           arm-build: ${{ matrix.config.platform == 'aarch64' }}
           manylinux: ${{ matrix.config.manylinux }}
       - uses: ./.github/workflows/upload_wheel
         with:
-          token: ${{ secrets.LANCEDB_PYPI_API_TOKEN }}
-          repo: "pypi"
+          pypi_token: ${{ secrets.LANCEDB_PYPI_API_TOKEN }}
+          fury_token: ${{ secrets.FURY_TOKEN }}
   mac:
-    # Only runs on tags that matches the python-make-release action
-    if: startsWith(github.ref, 'refs/tags/python-v')
     timeout-minutes: 60
     runs-on: ${{ matrix.config.runner }}
     strategy:
       matrix:
-        python-minor-version: ["8"]
         config:
           - target: x86_64-apple-darwin
             runner: macos-13
@@ -63,7 +58,6 @@ jobs:
     steps:
       - uses: actions/checkout@v4
         with:
-          ref: ${{ inputs.ref }}
           fetch-depth: 0
           lfs: true
       - name: Set up Python
@@ -72,38 +66,95 @@ jobs:
           python-version: 3.12
       - uses: ./.github/workflows/build_mac_wheel
         with:
-          python-minor-version: ${{ matrix.python-minor-version }}
+          python-minor-version: 8
           args: "--release --strip --target ${{ matrix.config.target }} --features fp16kernels"
       - uses: ./.github/workflows/upload_wheel
         with:
-          python-minor-version: ${{ matrix.python-minor-version }}
-          token: ${{ secrets.LANCEDB_PYPI_API_TOKEN }}
-          repo: "pypi"
+          pypi_token: ${{ secrets.LANCEDB_PYPI_API_TOKEN }}
+          fury_token: ${{ secrets.FURY_TOKEN }}
   windows:
-    # Only runs on tags that matches the python-make-release action
-    if: startsWith(github.ref, 'refs/tags/python-v')
     timeout-minutes: 60
     runs-on: windows-latest
-    strategy:
-      matrix:
-        python-minor-version: ["8"]
     steps:
       - uses: actions/checkout@v4
         with:
-          ref: ${{ inputs.ref }}
           fetch-depth: 0
           lfs: true
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: 3.${{ matrix.python-minor-version }}
+          python-version: 3.8
       - uses: ./.github/workflows/build_windows_wheel
         with:
-          python-minor-version: ${{ matrix.python-minor-version }}
+          python-minor-version: 8
           args: "--release --strip"
           vcpkg_token: ${{ secrets.VCPKG_GITHUB_PACKAGES }}
       - uses: ./.github/workflows/upload_wheel
         with:
-          python-minor-version: ${{ matrix.python-minor-version }}
-          token: ${{ secrets.LANCEDB_PYPI_API_TOKEN }}
-          repo: "pypi"
+          pypi_token: ${{ secrets.LANCEDB_PYPI_API_TOKEN }}
+          fury_token: ${{ secrets.FURY_TOKEN }}
+  gh-release:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Extract version
+        id: extract_version
+        env:
+          GITHUB_REF: ${{ github.ref }}
+        run: |
+          set -e
+          echo "Extracting tag and version from $GITHUB_REF"
+          if [[ $GITHUB_REF =~ refs/tags/python-v(.*) ]]; then
+            VERSION=${BASH_REMATCH[1]}
+            TAG=python-v$VERSION
+            echo "tag=$TAG" >> $GITHUB_OUTPUT
+            echo "version=$VERSION" >> $GITHUB_OUTPUT
+          else
+            echo "Failed to extract version from $GITHUB_REF"
+            exit 1
+          fi
+          echo "Extracted version $VERSION from $GITHUB_REF"
+          if [[ $VERSION =~ beta ]]; then
+            echo "This is a beta release"
+
+            # Get last release (that is not this one)
+            FROM_TAG=$(git tag --sort='version:refname' \
+              | grep ^python-v \
+              | grep -vF "$TAG" \
+              | python ci/semver_sort.py python-v \
+              | tail -n 1)
+          else
+            echo "This is a stable release"
+            # Get last stable tag (ignore betas)
+            FROM_TAG=$(git tag --sort='version:refname' \
+              | grep ^python-v \
+              | grep -vF "$TAG" \
+              | grep -v beta \
+              | python ci/semver_sort.py python-v \
+              | tail -n 1)
+          fi
+          echo "Found from tag $FROM_TAG"
+          echo "from_tag=$FROM_TAG" >> $GITHUB_OUTPUT
+      - name: Create Python Release Notes
+        id: python_release_notes
+        uses: mikepenz/release-changelog-builder-action@v4
+        with:
+          configuration: .github/release_notes.json
+          toTag: ${{ steps.extract_version.outputs.tag }}
+          fromTag: ${{ steps.extract_version.outputs.from_tag }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      - name: Create Python GH release
+        uses: softprops/action-gh-release@v2
+        with:
+          prerelease: ${{ contains('beta', github.ref) }}
+          tag_name: ${{ steps.extract_version.outputs.tag }}
+          token: ${{ secrets.GITHUB_TOKEN }}
+          generate_release_notes: false
+          name: Python LanceDB v${{ steps.extract_version.outputs.version }}
+          body: ${{ steps.python_release_notes.outputs.changelog }}

.github/workflows/python-make-release-commit.yml

@@ -1,56 +0,0 @@
-name: Python - Create release commit
-
-on:
-  workflow_dispatch:
-    inputs:
-      dry_run:
-        description: 'Dry run (create the local commit/tags but do not push it)'
-        required: true
-        default: "false"
-        type: choice
-        options:
-          - "true"
-          - "false"
-      part:
-        description: 'What kind of release is this?'
-        required: true
-        default: 'patch'
-        type: choice
-        options:
-          - patch
-          - minor
-          - major
-
-jobs:
-  bump-version:
-    runs-on: ubuntu-latest
-    steps:
-    - name: Check out main
-      uses: actions/checkout@v4
-      with:
-        ref: main
-        persist-credentials: false
-        fetch-depth: 0
-        lfs: true
-    - name: Set git configs for bumpversion
-      shell: bash
-      run: |
-        git config user.name 'Lance Release'
-        git config user.email 'lance-dev@lancedb.com'
-    - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: "3.11"
-    - name: Bump version, create tag and commit
-      working-directory: python
-      run: |
-        pip install bump2version
-        bumpversion --verbose ${{ inputs.part }}
-    - name: Push new version and tag
-      if: ${{ inputs.dry_run }} == "false"
-      uses: ad-m/github-push-action@master
-      with:
-        github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
-        branch: main
-        tags: true
-

.github/workflows/python.yml

@@ -75,7 +75,7 @@ jobs:
     timeout-minutes: 30
     strategy:
       matrix:
-        python-minor-version: ["8", "11"]
+        python-minor-version: ["9", "11"]
     runs-on: "ubuntu-22.04"
     defaults:
       run:

.github/workflows/rust.yml

@@ -74,11 +74,11 @@ jobs:
       run: |
           sudo apt update
           sudo apt install -y protobuf-compiler libssl-dev
-    - name: Build
-      run: cargo build --all-features
     - name: Start S3 integration test environment
       working-directory: .
       run: docker compose up --detach --wait
+    - name: Build
+      run: cargo build --all-features
     - name: Run tests
       run: cargo test --all-features
     - name: Run examples

.github/workflows/upload_wheel/action.yml

@@ -2,28 +2,43 @@ name: upload-wheel
 
 description: "Upload wheels to Pypi"
 inputs:
-  os:
-    required: true
-    description: "ubuntu-22.04 or macos-13"
-  repo:
-    required: false
-    description: "pypi or testpypi"
-    default: "pypi"
-  token:
+  pypi_token:
     required: true
     description: "release token for the repo"
+  fury_token:
+    required: true
+    description: "release token for the fury repo"
 
 runs:
   using: "composite"
   steps:
-    - name: Install dependencies
-      shell: bash
-      run: |
-        python -m pip install --upgrade pip
-        pip install twine
-    - name: Publish wheel
-      env:
-        TWINE_USERNAME: __token__
-        TWINE_PASSWORD: ${{ inputs.token }}
-      shell: bash
-      run: twine upload --repository ${{ inputs.repo }} target/wheels/lancedb-*.whl
+  - name: Install dependencies
+    shell: bash
+    run: |
+      python -m pip install --upgrade pip
+      pip install twine
+  - name: Choose repo
+    shell: bash
+    id: choose_repo
+    run: |
+      if [ ${{ github.ref }} == "*beta*" ]; then
+        echo "repo=fury" >> $GITHUB_OUTPUT
+      else
+        echo "repo=pypi" >> $GITHUB_OUTPUT
+      fi
+  - name: Publish to PyPI
+    shell: bash
+    env:
+      FURY_TOKEN: ${{ inputs.fury_token }}
+      PYPI_TOKEN: ${{ inputs.pypi_token }}
+    run: |
+      if [ ${{ steps.choose_repo.outputs.repo }} == "fury" ]; then
+        WHEEL=$(ls target/wheels/lancedb-*.whl 2> /dev/null | head -n 1)
+        echo "Uploading $WHEEL to Fury"
+        curl -f -F package=@$WHEEL https://$FURY_TOKEN@push.fury.io/lancedb/
+      else
+        twine upload --repository ${{ steps.choose_repo.outputs.repo }} \
+          --username __token__ \
+          --password $PYPI_TOKEN \
+          target/wheels/lancedb-*.whl
+      fi

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
   hooks:
     - id: local-biome-check
       name: biome check
-      entry: npx biome check
+      entry: npx @biomejs/biome check --config-path nodejs/biome.json nodejs/
       language: system
       types: [text]
       files: "nodejs/.*"

Cargo.toml

@@ -1,5 +1,11 @@
 [workspace]
-members = ["rust/ffi/node", "rust/lancedb", "nodejs", "python"]
+members = [
+    "rust/ffi/node",
+    "rust/lancedb",
+    "nodejs",
+    "python",
+    "java/core/lancedb-jni",
+]
 # Python package needs to be built by maturin.
 exclude = ["python"]
 resolver = "2"
@@ -14,10 +20,11 @@ keywords = ["lancedb", "lance", "database", "vector", "search"]
 categories = ["database-implementations"]
 
 [workspace.dependencies]
-lance = { "version" = "=0.10.18", "features" = ["dynamodb"] }
-lance-index = { "version" = "=0.10.18" }
-lance-linalg = { "version" = "=0.10.18" }
-lance-testing = { "version" = "=0.10.18" }
+lance = { "version" = "=0.12.1", "features" = ["dynamodb"] }
+lance-index = { "version" = "=0.12.1" }
+lance-linalg = { "version" = "=0.12.1" }
+lance-testing = { "version" = "=0.12.1" }
+lance-datafusion = { "version" = "=0.12.1" }
 # Note that this one does not include pyarrow
 arrow = { version = "51.0", optional = false }
 arrow-array = "51.0"
@@ -29,6 +36,7 @@ arrow-arith = "51.0"
 arrow-cast = "51.0"
 async-trait = "0"
 chrono = "0.4.35"
+datafusion-physical-plan = "37.1"
 half = { "version" = "=2.4.1", default-features = false, features = [
     "num-traits",
 ] }

README.md

@@ -83,5 +83,5 @@ result = table.search([100, 100]).limit(2).to_pandas()
 ```
 
 ## Blogs, Tutorials & Videos
-* 📈 <a href="https://blog.eto.ai/benchmarking-random-access-in-lance-ed690757a826">2000x better performance with Lance over Parquet</a>
+* 📈 <a href="https://blog.lancedb.com/benchmarking-random-access-in-lance/">2000x better performance with Lance over Parquet</a>
 * 🤖 <a href="https://github.com/lancedb/lancedb/blob/main/docs/src/notebooks/youtube_transcript_search.ipynb">Build a question and answer bot with LanceDB</a>

ci/bump_version.sh

@@ -0,0 +1,51 @@
+set -e
+
+RELEASE_TYPE=${1:-"stable"}
+BUMP_MINOR=${2:-false}
+TAG_PREFIX=${3:-"v"} # Such as "python-v"
+HEAD_SHA=${4:-$(git rev-parse HEAD)}
+
+readonly SELF_DIR=$(cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )
+
+PREV_TAG=$(git tag --sort='version:refname' | grep ^$TAG_PREFIX | python $SELF_DIR/semver_sort.py $TAG_PREFIX | tail -n 1)
+echo "Found previous tag $PREV_TAG"
+
+# Initially, we don't want to tag if we are doing stable, because we will bump
+# again later. See comment at end for why.
+if [[ "$RELEASE_TYPE" == 'stable' ]]; then 
+  BUMP_ARGS="--no-tag"
+fi
+
+# If last is stable and not bumping minor
+if [[ $PREV_TAG != *beta* ]]; then
+    if [[ "$BUMP_MINOR" != "false" ]]; then
+      # X.Y.Z -> X.(Y+1).0-beta.0
+      bump-my-version bump -vv $BUMP_ARGS minor
+    else
+      # X.Y.Z -> X.Y.(Z+1)-beta.0
+      bump-my-version bump -vv $BUMP_ARGS patch
+    fi
+else
+  if [[ "$BUMP_MINOR" != "false" ]]; then
+    # X.Y.Z-beta.N -> X.(Y+1).0-beta.0
+    bump-my-version bump -vv $BUMP_ARGS minor
+  else
+    # X.Y.Z-beta.N -> X.Y.Z-beta.(N+1)
+    bump-my-version bump -vv $BUMP_ARGS pre_n
+  fi
+fi
+
+# The above bump will always bump to a pre-release version. If we are releasing
+# a stable version, bump the pre-release level ("pre_l") to make it stable.
+if [[ $RELEASE_TYPE == 'stable' ]]; then
+  # X.Y.Z-beta.N -> X.Y.Z
+  bump-my-version bump -vv pre_l
+fi
+
+# Validate that we have incremented version appropriately for breaking changes
+NEW_TAG=$(git describe --tags --exact-match HEAD)
+NEW_VERSION=$(echo $NEW_TAG | sed "s/^$TAG_PREFIX//")
+LAST_STABLE_RELEASE=$(git tag --sort='version:refname' | grep ^$TAG_PREFIX | grep -v beta | grep -vF "$NEW_TAG" | python $SELF_DIR/semver_sort.py $TAG_PREFIX | tail -n 1)
+LAST_STABLE_VERSION=$(echo $LAST_STABLE_RELEASE | sed "s/^$TAG_PREFIX//")
+
+python $SELF_DIR/check_breaking_changes.py $LAST_STABLE_RELEASE $HEAD_SHA $LAST_STABLE_VERSION $NEW_VERSION

ci/check_breaking_changes.py

@@ -0,0 +1,35 @@
+"""
+Check whether there are any breaking changes in the PRs between the base and head commits.
+If there are, assert that we have incremented the minor version.
+"""
+import argparse
+import os
+from packaging.version import parse
+
+from github import Github
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("base")
+    parser.add_argument("head")
+    parser.add_argument("last_stable_version")
+    parser.add_argument("current_version")
+    args = parser.parse_args()
+
+    repo = Github(os.environ["GITHUB_TOKEN"]).get_repo(os.environ["GITHUB_REPOSITORY"])
+    commits = repo.compare(args.base, args.head).commits
+    prs = (pr for commit in commits for pr in commit.get_pulls())
+
+    for pr in prs:
+        if any(label.name == "breaking-change" for label in pr.labels):
+            print(f"Breaking change in PR: {pr.html_url}")
+            break
+    else:
+        print("No breaking changes found.")
+        exit(0)
+    
+    last_stable_version = parse(args.last_stable_version)
+    current_version = parse(args.current_version)
+    if current_version.minor <= last_stable_version.minor:
+        print("Minor version is not greater than the last stable version.")
+        exit(1)

ci/semver_sort.py

@@ -0,0 +1,35 @@
+"""
+Takes a list of semver strings and sorts them in ascending order.
+"""
+
+import sys
+from packaging.version import parse, InvalidVersion
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("prefix", default="v")
+    args = parser.parse_args()
+
+    # Read the input from stdin
+    lines = sys.stdin.readlines()
+
+    # Parse the versions
+    versions = []
+    for line in lines:
+        line = line.strip()
+        try:
+            version_str = line.removeprefix(args.prefix)
+            version = parse(version_str)
+        except InvalidVersion:
+            # There are old tags that don't follow the semver format
+            print(f"Invalid version: {line}", file=sys.stderr)
+            continue
+        versions.append((line, version))
+
+    # Sort the versions
+    versions.sort(key=lambda x: x[1])
+
+    # Print the sorted versions as original strings
+    for line, _ in versions:
+        print(line)

docs/src/basic.md

@@ -44,6 +44,36 @@
 
     !!! info "Please also make sure you're using the same version of Arrow as in the [lancedb crate](https://github.com/lancedb/lancedb/blob/main/Cargo.toml)"
 
+### Preview releases
+
+Stable releases are created about every 2 weeks. For the latest features and bug
+fixes, you can install the preview release. These releases receive the same
+level of testing as stable releases, but are not guaranteed to be available for
+more than 6 months after they are released. Once your application is stable, we
+recommend switching to stable releases.
+
+=== "Python"
+
+      ```shell
+      pip install --pre --extra-index-url https://pypi.fury.io/lancedb/ lancedb
+      ```
+
+=== "Typescript"
+
+      ```shell
+      npm install vectordb@preview
+      ```
+
+=== "Rust"
+    
+    We don't push preview releases to crates.io, but you can referent the tag
+    in GitHub within your Cargo dependencies:
+
+    ```toml
+    [dependencies]
+    lancedb = { git = "https://github.com/lancedb/lancedb.git", tag = "vX.Y.Z-beta.N" }
+    ```
+
 ## Connect to a database
 
 === "Python"

java/core/lancedb-jni/Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "lancedb-jni"
+description = "JNI bindings for LanceDB"
+# TODO modify lancedb/Cargo.toml for version and dependencies
+version = "0.4.18"
+edition.workspace = true
+repository.workspace = true
+readme.workspace = true
+license.workspace = true
+keywords.workspace = true
+categories.workspace = true
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+lancedb = { path = "../../../rust/lancedb" }
+lance = { workspace = true }
+arrow = { workspace = true, features = ["ffi"] }
+arrow-schema.workspace = true
+tokio = "1.23"
+jni = "0.21.1"
+snafu.workspace = true
+lazy_static.workspace = true
+serde = { version = "^1" }
+serde_json = { version = "1" }

java/core/lancedb-jni/src/connection.rs

@@ -0,0 +1,130 @@
+use crate::ffi::JNIEnvExt;
+use crate::traits::IntoJava;
+use crate::{Error, RT};
+use jni::objects::{JObject, JString, JValue};
+use jni::JNIEnv;
+pub const NATIVE_CONNECTION: &str = "nativeConnectionHandle";
+use crate::Result;
+use lancedb::connection::{connect, Connection};
+
+#[derive(Clone)]
+pub struct BlockingConnection {
+    pub(crate) inner: Connection,
+}
+
+impl BlockingConnection {
+    pub fn create(dataset_uri: &str) -> Result<Self> {
+        let inner = RT.block_on(connect(dataset_uri).execute())?;
+        Ok(Self { inner })
+    }
+
+    pub fn table_names(
+        &self,
+        start_after: Option<String>,
+        limit: Option<i32>,
+    ) -> Result<Vec<String>> {
+        let mut op = self.inner.table_names();
+        if let Some(start_after) = start_after {
+            op = op.start_after(start_after);
+        }
+        if let Some(limit) = limit {
+            op = op.limit(limit as u32);
+        }
+        Ok(RT.block_on(op.execute())?)
+    }
+}
+
+impl IntoJava for BlockingConnection {
+    fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> JObject<'a> {
+        attach_native_connection(env, self)
+    }
+}
+
+fn attach_native_connection<'local>(
+    env: &mut JNIEnv<'local>,
+    connection: BlockingConnection,
+) -> JObject<'local> {
+    let j_connection = create_java_connection_object(env);
+    // This block sets a native Rust object (Connection) as a field in the Java object (j_Connection).
+    // Caution: This creates a potential for memory leaks. The Rust object (Connection) is not
+    // automatically garbage-collected by Java, and its memory will not be freed unless
+    // explicitly handled.
+    //
+    // To prevent memory leaks, ensure the following:
+    // 1. The Java object (`j_Connection`) should implement the `java.io.Closeable` interface.
+    // 2. Users of this Java object should be instructed to always use it within a try-with-resources
+    //    statement (or manually call the `close()` method) to ensure that `self.close()` is invoked.
+    match unsafe { env.set_rust_field(&j_connection, NATIVE_CONNECTION, connection) } {
+        Ok(_) => j_connection,
+        Err(err) => {
+            env.throw_new(
+                "java/lang/RuntimeException",
+                format!("Failed to set native handle for Connection: {}", err),
+            )
+            .expect("Error throwing exception");
+            JObject::null()
+        }
+    }
+}
+
+fn create_java_connection_object<'a>(env: &mut JNIEnv<'a>) -> JObject<'a> {
+    env.new_object("com/lancedb/lancedb/Connection", "()V", &[])
+        .expect("Failed to create Java Lance Connection instance")
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lancedb_Connection_releaseNativeConnection(
+    mut env: JNIEnv,
+    j_connection: JObject,
+) {
+    let _: BlockingConnection = unsafe {
+        env.take_rust_field(j_connection, NATIVE_CONNECTION)
+            .expect("Failed to take native Connection handle")
+    };
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lancedb_Connection_connect<'local>(
+    mut env: JNIEnv<'local>,
+    _obj: JObject,
+    dataset_uri_object: JString,
+) -> JObject<'local> {
+    let dataset_uri: String = ok_or_throw!(env, env.get_string(&dataset_uri_object)).into();
+    let blocking_connection = ok_or_throw!(env, BlockingConnection::create(&dataset_uri));
+    blocking_connection.into_java(&mut env)
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lancedb_Connection_tableNames<'local>(
+    mut env: JNIEnv<'local>,
+    j_connection: JObject,
+    start_after_obj: JObject, // Optional<String>
+    limit_obj: JObject,       // Optional<Integer>
+) -> JObject<'local> {
+    ok_or_throw!(
+        env,
+        inner_table_names(&mut env, j_connection, start_after_obj, limit_obj)
+    )
+}
+
+fn inner_table_names<'local>(
+    env: &mut JNIEnv<'local>,
+    j_connection: JObject,
+    start_after_obj: JObject, // Optional<String>
+    limit_obj: JObject,       // Optional<Integer>
+) -> Result<JObject<'local>> {
+    let start_after = env.get_string_opt(&start_after_obj)?;
+    let limit = env.get_int_opt(&limit_obj)?;
+    let conn =
+        unsafe { env.get_rust_field::<_, _, BlockingConnection>(j_connection, NATIVE_CONNECTION) }?;
+    let table_names = conn.table_names(start_after, limit)?;
+    drop(conn);
+    let j_names = env.new_object("java/util/ArrayList", "()V", &[])?;
+    for item in table_names {
+        let jstr_item = env.new_string(item)?;
+        let item_jobj = JObject::from(jstr_item);
+        let item_gen = JValue::Object(&item_jobj);
+        env.call_method(&j_names, "add", "(Ljava/lang/Object;)Z", &[item_gen])?;
+    }
+    Ok(j_names)
+}

java/core/lancedb-jni/src/error.rs

@@ -0,0 +1,225 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::str::Utf8Error;
+
+use arrow_schema::ArrowError;
+use jni::errors::Error as JniError;
+use serde_json::Error as JsonError;
+use snafu::{Location, Snafu};
+
+type BoxedError = Box<dyn std::error::Error + Send + Sync + 'static>;
+
+/// Java Exception types
+pub enum JavaException {
+    IllegalArgumentException,
+    IOException,
+    RuntimeException,
+}
+
+impl JavaException {
+    pub fn as_str(&self) -> &str {
+        match self {
+            Self::IllegalArgumentException => "java/lang/IllegalArgumentException",
+            Self::IOException => "java/io/IOException",
+            Self::RuntimeException => "java/lang/RuntimeException",
+        }
+    }
+}
+/// TODO(lu) change to lancedb-jni
+#[derive(Debug, Snafu)]
+#[snafu(visibility(pub))]
+pub enum Error {
+    #[snafu(display("JNI error: {message}, {location}"))]
+    Jni { message: String, location: Location },
+    #[snafu(display("Invalid argument: {message}, {location}"))]
+    InvalidArgument { message: String, location: Location },
+    #[snafu(display("IO error: {source}, {location}"))]
+    IO {
+        source: BoxedError,
+        location: Location,
+    },
+    #[snafu(display("Arrow error: {message}, {location}"))]
+    Arrow { message: String, location: Location },
+    #[snafu(display("Index error: {message}, {location}"))]
+    Index { message: String, location: Location },
+    #[snafu(display("JSON error: {message}, {location}"))]
+    JSON { message: String, location: Location },
+    #[snafu(display("Dataset at path {path} was not found, {location}"))]
+    DatasetNotFound { path: String, location: Location },
+    #[snafu(display("Dataset already exists: {uri}, {location}"))]
+    DatasetAlreadyExists { uri: String, location: Location },
+    #[snafu(display("Table '{name}' already exists"))]
+    TableAlreadyExists { name: String },
+    #[snafu(display("Table '{name}' was not found"))]
+    TableNotFound { name: String },
+    #[snafu(display("Invalid table name '{name}': {reason}"))]
+    InvalidTableName { name: String, reason: String },
+    #[snafu(display("Embedding function '{name}' was not found: {reason}, {location}"))]
+    EmbeddingFunctionNotFound {
+        name: String,
+        reason: String,
+        location: Location,
+    },
+    #[snafu(display("Other Lance error: {message}, {location}"))]
+    OtherLance { message: String, location: Location },
+    #[snafu(display("Other LanceDB error: {message}, {location}"))]
+    OtherLanceDB { message: String, location: Location },
+}
+
+impl Error {
+    /// Throw as Java Exception
+    pub fn throw(&self, env: &mut jni::JNIEnv) {
+        match self {
+            Self::InvalidArgument { .. }
+            | Self::DatasetNotFound { .. }
+            | Self::DatasetAlreadyExists { .. }
+            | Self::TableAlreadyExists { .. }
+            | Self::TableNotFound { .. }
+            | Self::InvalidTableName { .. }
+            | Self::EmbeddingFunctionNotFound { .. } => {
+                self.throw_as(env, JavaException::IllegalArgumentException)
+            }
+            Self::IO { .. } | Self::Index { .. } => self.throw_as(env, JavaException::IOException),
+            Self::Arrow { .. }
+            | Self::JSON { .. }
+            | Self::OtherLance { .. }
+            | Self::OtherLanceDB { .. }
+            | Self::Jni { .. } => self.throw_as(env, JavaException::RuntimeException),
+        }
+    }
+
+    /// Throw as an concrete Java Exception
+    pub fn throw_as(&self, env: &mut jni::JNIEnv, exception: JavaException) {
+        let message = &format!(
+            "Error when throwing Java exception: {}:{}",
+            exception.as_str(),
+            self
+        );
+        env.throw_new(exception.as_str(), self.to_string())
+            .expect(message);
+    }
+}
+
+pub type Result<T> = std::result::Result<T, Error>;
+
+trait ToSnafuLocation {
+    fn to_snafu_location(&'static self) -> snafu::Location;
+}
+
+impl ToSnafuLocation for std::panic::Location<'static> {
+    fn to_snafu_location(&'static self) -> snafu::Location {
+        snafu::Location::new(self.file(), self.line(), self.column())
+    }
+}
+
+impl From<JniError> for Error {
+    #[track_caller]
+    fn from(source: JniError) -> Self {
+        Self::Jni {
+            message: source.to_string(),
+            location: std::panic::Location::caller().to_snafu_location(),
+        }
+    }
+}
+
+impl From<Utf8Error> for Error {
+    #[track_caller]
+    fn from(source: Utf8Error) -> Self {
+        Self::InvalidArgument {
+            message: source.to_string(),
+            location: std::panic::Location::caller().to_snafu_location(),
+        }
+    }
+}
+
+impl From<ArrowError> for Error {
+    #[track_caller]
+    fn from(source: ArrowError) -> Self {
+        Self::Arrow {
+            message: source.to_string(),
+            location: std::panic::Location::caller().to_snafu_location(),
+        }
+    }
+}
+
+impl From<JsonError> for Error {
+    #[track_caller]
+    fn from(source: JsonError) -> Self {
+        Self::JSON {
+            message: source.to_string(),
+            location: std::panic::Location::caller().to_snafu_location(),
+        }
+    }
+}
+
+impl From<lance::Error> for Error {
+    #[track_caller]
+    fn from(source: lance::Error) -> Self {
+        match source {
+            lance::Error::DatasetNotFound {
+                path,
+                source: _,
+                location,
+            } => Self::DatasetNotFound { path, location },
+            lance::Error::DatasetAlreadyExists { uri, location } => {
+                Self::DatasetAlreadyExists { uri, location }
+            }
+            lance::Error::IO { source, location } => Self::IO { source, location },
+            lance::Error::Arrow { message, location } => Self::Arrow { message, location },
+            lance::Error::Index { message, location } => Self::Index { message, location },
+            lance::Error::InvalidInput { source, location } => Self::InvalidArgument {
+                message: source.to_string(),
+                location,
+            },
+            _ => Self::OtherLance {
+                message: source.to_string(),
+                location: std::panic::Location::caller().to_snafu_location(),
+            },
+        }
+    }
+}
+
+impl From<lancedb::Error> for Error {
+    #[track_caller]
+    fn from(source: lancedb::Error) -> Self {
+        match source {
+            lancedb::Error::InvalidTableName { name, reason } => {
+                Self::InvalidTableName { name, reason }
+            }
+            lancedb::Error::InvalidInput { message } => Self::InvalidArgument {
+                message,
+                location: std::panic::Location::caller().to_snafu_location(),
+            },
+            lancedb::Error::TableNotFound { name } => Self::TableNotFound { name },
+            lancedb::Error::TableAlreadyExists { name } => Self::TableAlreadyExists { name },
+            lancedb::Error::EmbeddingFunctionNotFound { name, reason } => {
+                Self::EmbeddingFunctionNotFound {
+                    name,
+                    reason,
+                    location: std::panic::Location::caller().to_snafu_location(),
+                }
+            }
+            lancedb::Error::Arrow { source } => Self::Arrow {
+                message: source.to_string(),
+                location: std::panic::Location::caller().to_snafu_location(),
+            },
+            lancedb::Error::Lance { source } => Self::from(source),
+            _ => Self::OtherLanceDB {
+                message: source.to_string(),
+                location: std::panic::Location::caller().to_snafu_location(),
+            },
+        }
+    }
+}

java/core/lancedb-jni/src/ffi.rs

@@ -0,0 +1,204 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use core::slice;
+
+use jni::objects::{JByteBuffer, JObjectArray, JString};
+use jni::sys::jobjectArray;
+use jni::{objects::JObject, JNIEnv};
+
+use crate::error::{Error, Result};
+
+/// TODO(lu) import from lance-jni without duplicate
+/// Extend JNIEnv with helper functions.
+pub trait JNIEnvExt {
+    /// Get integers from Java List<Integer> object.
+    fn get_integers(&mut self, obj: &JObject) -> Result<Vec<i32>>;
+
+    /// Get strings from Java List<String> object.
+    fn get_strings(&mut self, obj: &JObject) -> Result<Vec<String>>;
+
+    /// Get strings from Java String[] object.
+    /// Note that get Option<Vec<String>> from Java Optional<String[]> just doesn't work.
+    #[allow(unused)]
+    fn get_strings_array(&mut self, obj: jobjectArray) -> Result<Vec<String>>;
+
+    /// Get Option<String> from Java Optional<String>.
+    fn get_string_opt(&mut self, obj: &JObject) -> Result<Option<String>>;
+
+    /// Get Option<Vec<String>> from Java Optional<List<String>>.
+    #[allow(unused)]
+    fn get_strings_opt(&mut self, obj: &JObject) -> Result<Option<Vec<String>>>;
+
+    /// Get Option<i32> from Java Optional<Integer>.
+    fn get_int_opt(&mut self, obj: &JObject) -> Result<Option<i32>>;
+
+    /// Get Option<Vec<i32>> from Java Optional<List<Integer>>.
+    fn get_ints_opt(&mut self, obj: &JObject) -> Result<Option<Vec<i32>>>;
+
+    /// Get Option<i64> from Java Optional<Long>.
+    #[allow(unused)]
+    fn get_long_opt(&mut self, obj: &JObject) -> Result<Option<i64>>;
+
+    /// Get Option<u64> from Java Optional<Long>.
+    #[allow(unused)]
+    fn get_u64_opt(&mut self, obj: &JObject) -> Result<Option<u64>>;
+
+    /// Get Option<&[u8]> from Java Optional<ByteBuffer>.
+    #[allow(unused)]
+    fn get_bytes_opt(&mut self, obj: &JObject) -> Result<Option<&[u8]>>;
+
+    fn get_optional<T, F>(&mut self, obj: &JObject, f: F) -> Result<Option<T>>
+    where
+        F: FnOnce(&mut JNIEnv, &JObject) -> Result<T>;
+}
+
+impl JNIEnvExt for JNIEnv<'_> {
+    fn get_integers(&mut self, obj: &JObject) -> Result<Vec<i32>> {
+        let list = self.get_list(obj)?;
+        let mut iter = list.iter(self)?;
+        let mut results = Vec::with_capacity(list.size(self)? as usize);
+        while let Some(elem) = iter.next(self)? {
+            let int_obj = self.call_method(elem, "intValue", "()I", &[])?;
+            let int_value = int_obj.i()?;
+            results.push(int_value);
+        }
+        Ok(results)
+    }
+
+    fn get_strings(&mut self, obj: &JObject) -> Result<Vec<String>> {
+        let list = self.get_list(obj)?;
+        let mut iter = list.iter(self)?;
+        let mut results = Vec::with_capacity(list.size(self)? as usize);
+        while let Some(elem) = iter.next(self)? {
+            let jstr = JString::from(elem);
+            let val = self.get_string(&jstr)?;
+            results.push(val.to_str()?.to_string())
+        }
+        Ok(results)
+    }
+
+    fn get_strings_array(&mut self, obj: jobjectArray) -> Result<Vec<String>> {
+        let jobject_array = unsafe { JObjectArray::from_raw(obj) };
+        let array_len = self.get_array_length(&jobject_array)?;
+        let mut res: Vec<String> = Vec::new();
+        for i in 0..array_len {
+            let item: JString = self.get_object_array_element(&jobject_array, i)?.into();
+            res.push(self.get_string(&item)?.into());
+        }
+        Ok(res)
+    }
+
+    fn get_string_opt(&mut self, obj: &JObject) -> Result<Option<String>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_string_obj = java_obj_gen.l()?;
+            let jstr = JString::from(java_string_obj);
+            let val = env.get_string(&jstr)?;
+            Ok(val.to_str()?.to_string())
+        })
+    }
+
+    fn get_strings_opt(&mut self, obj: &JObject) -> Result<Option<Vec<String>>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_list_obj = java_obj_gen.l()?;
+            env.get_strings(&java_list_obj)
+        })
+    }
+
+    fn get_int_opt(&mut self, obj: &JObject) -> Result<Option<i32>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_int_obj = java_obj_gen.l()?;
+            let int_obj = env.call_method(java_int_obj, "intValue", "()I", &[])?;
+            let int_value = int_obj.i()?;
+            Ok(int_value)
+        })
+    }
+
+    fn get_ints_opt(&mut self, obj: &JObject) -> Result<Option<Vec<i32>>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_list_obj = java_obj_gen.l()?;
+            env.get_integers(&java_list_obj)
+        })
+    }
+
+    fn get_long_opt(&mut self, obj: &JObject) -> Result<Option<i64>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_long_obj = java_obj_gen.l()?;
+            let long_obj = env.call_method(java_long_obj, "longValue", "()J", &[])?;
+            let long_value = long_obj.j()?;
+            Ok(long_value)
+        })
+    }
+
+    fn get_u64_opt(&mut self, obj: &JObject) -> Result<Option<u64>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_long_obj = java_obj_gen.l()?;
+            let long_obj = env.call_method(java_long_obj, "longValue", "()J", &[])?;
+            let long_value = long_obj.j()?;
+            Ok(long_value as u64)
+        })
+    }
+
+    fn get_bytes_opt(&mut self, obj: &JObject) -> Result<Option<&[u8]>> {
+        self.get_optional(obj, |env, inner_obj| {
+            let java_obj_gen = env.call_method(inner_obj, "get", "()Ljava/lang/Object;", &[])?;
+            let java_byte_buffer_obj = java_obj_gen.l()?;
+            let j_byte_buffer = JByteBuffer::from(java_byte_buffer_obj);
+            let raw_data = env.get_direct_buffer_address(&j_byte_buffer)?;
+            let capacity = env.get_direct_buffer_capacity(&j_byte_buffer)?;
+            let data = unsafe { slice::from_raw_parts(raw_data, capacity) };
+            Ok(data)
+        })
+    }
+
+    fn get_optional<T, F>(&mut self, obj: &JObject, f: F) -> Result<Option<T>>
+    where
+        F: FnOnce(&mut JNIEnv, &JObject) -> Result<T>,
+    {
+        if obj.is_null() {
+            return Ok(None);
+        }
+        let is_empty = self.call_method(obj, "isEmpty", "()Z", &[])?;
+        if is_empty.z()? {
+            // TODO(lu): put get java object into here cuz can only get java Object
+            Ok(None)
+        } else {
+            f(self, obj).map(Some)
+        }
+    }
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_test_JniTestHelper_parseInts(
+    mut env: JNIEnv,
+    _obj: JObject,
+    list_obj: JObject, // List<Integer>
+) {
+    ok_or_throw_without_return!(env, env.get_integers(&list_obj));
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_test_JniTestHelper_parseIntsOpt(
+    mut env: JNIEnv,
+    _obj: JObject,
+    list_obj: JObject, // Optional<List<Integer>>
+) {
+    ok_or_throw_without_return!(env, env.get_ints_opt(&list_obj));
+}

java/core/lancedb-jni/src/lib.rs

@@ -0,0 +1,68 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use lazy_static::lazy_static;
+
+// TODO import from lance-jni without duplicate
+#[macro_export]
+macro_rules! ok_or_throw {
+    ($env:expr, $result:expr) => {
+        match $result {
+            Ok(value) => value,
+            Err(err) => {
+                Error::from(err).throw(&mut $env);
+                return JObject::null();
+            }
+        }
+    };
+}
+
+macro_rules! ok_or_throw_without_return {
+    ($env:expr, $result:expr) => {
+        match $result {
+            Ok(value) => value,
+            Err(err) => {
+                Error::from(err).throw(&mut $env);
+                return;
+            }
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! ok_or_throw_with_return {
+    ($env:expr, $result:expr, $ret:expr) => {
+        match $result {
+            Ok(value) => value,
+            Err(err) => {
+                Error::from(err).throw(&mut $env);
+                return $ret;
+            }
+        }
+    };
+}
+
+mod connection;
+pub mod error;
+mod ffi;
+mod traits;
+
+pub use error::{Error, Result};
+
+lazy_static! {
+    static ref RT: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .expect("Failed to create tokio runtime");
+}

java/core/lancedb-jni/src/traits.rs

@@ -0,0 +1,122 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use jni::objects::{JMap, JObject, JString, JValue};
+use jni::JNIEnv;
+
+use crate::Result;
+
+pub trait FromJObject<T> {
+    fn extract(&self) -> Result<T>;
+}
+
+/// Convert a Rust type into a Java Object.
+pub trait IntoJava {
+    fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> JObject<'a>;
+}
+
+impl FromJObject<i32> for JObject<'_> {
+    fn extract(&self) -> Result<i32> {
+        Ok(JValue::from(self).i()?)
+    }
+}
+
+impl FromJObject<i64> for JObject<'_> {
+    fn extract(&self) -> Result<i64> {
+        Ok(JValue::from(self).j()?)
+    }
+}
+
+impl FromJObject<f32> for JObject<'_> {
+    fn extract(&self) -> Result<f32> {
+        Ok(JValue::from(self).f()?)
+    }
+}
+
+impl FromJObject<f64> for JObject<'_> {
+    fn extract(&self) -> Result<f64> {
+        Ok(JValue::from(self).d()?)
+    }
+}
+
+pub trait FromJString {
+    fn extract(&self, env: &mut JNIEnv) -> Result<String>;
+}
+
+impl FromJString for JString<'_> {
+    fn extract(&self, env: &mut JNIEnv) -> Result<String> {
+        Ok(env.get_string(self)?.into())
+    }
+}
+
+pub trait JMapExt {
+    #[allow(dead_code)]
+    fn get_string(&self, env: &mut JNIEnv, key: &str) -> Result<Option<String>>;
+
+    #[allow(dead_code)]
+    fn get_i32(&self, env: &mut JNIEnv, key: &str) -> Result<Option<i32>>;
+
+    #[allow(dead_code)]
+    fn get_i64(&self, env: &mut JNIEnv, key: &str) -> Result<Option<i64>>;
+
+    #[allow(dead_code)]
+    fn get_f32(&self, env: &mut JNIEnv, key: &str) -> Result<Option<f32>>;
+
+    #[allow(dead_code)]
+    fn get_f64(&self, env: &mut JNIEnv, key: &str) -> Result<Option<f64>>;
+}
+
+fn get_map_value<T>(env: &mut JNIEnv, map: &JMap, key: &str) -> Result<Option<T>>
+where
+    for<'a> JObject<'a>: FromJObject<T>,
+{
+    let key_obj: JObject = env.new_string(key)?.into();
+    if let Some(value) = map.get(env, &key_obj)? {
+        if value.is_null() {
+            Ok(None)
+        } else {
+            Ok(Some(value.extract()?))
+        }
+    } else {
+        Ok(None)
+    }
+}
+
+impl JMapExt for JMap<'_, '_, '_> {
+    fn get_string(&self, env: &mut JNIEnv, key: &str) -> Result<Option<String>> {
+        let key_obj: JObject = env.new_string(key)?.into();
+        if let Some(value) = self.get(env, &key_obj)? {
+            let value_str: JString = value.into();
+            Ok(Some(value_str.extract(env)?))
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn get_i32(&self, env: &mut JNIEnv, key: &str) -> Result<Option<i32>> {
+        get_map_value(env, self, key)
+    }
+
+    fn get_i64(&self, env: &mut JNIEnv, key: &str) -> Result<Option<i64>> {
+        get_map_value(env, self, key)
+    }
+
+    fn get_f32(&self, env: &mut JNIEnv, key: &str) -> Result<Option<f32>> {
+        get_map_value(env, self, key)
+    }
+
+    fn get_f64(&self, env: &mut JNIEnv, key: &str) -> Result<Option<f64>> {
+        get_map_value(env, self, key)
+    }
+}

java/core/pom.xml

@@ -0,0 +1,94 @@
+<?xml version="1.0" encoding="UTF-8"?>
+
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>com.lancedb</groupId>
+        <artifactId>lancedb-parent</artifactId>
+        <version>0.1-SNAPSHOT</version>
+        <relativePath>../pom.xml</relativePath>
+    </parent>
+
+    <artifactId>lancedb-core</artifactId>
+    <name>LanceDB Core</name>
+    <packaging>jar</packaging>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-vector</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-memory-netty</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-c-data</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-dataset</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.json</groupId>
+            <artifactId>json</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.questdb</groupId>
+            <artifactId>jar-jni</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter</artifactId>
+             <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <profiles>
+        <profile>
+            <id>build-jni</id>
+            <activation>
+                <activeByDefault>true</activeByDefault>
+            </activation>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.questdb</groupId>
+                        <artifactId>rust-maven-plugin</artifactId>
+                        <version>1.1.1</version>
+                        <executions>
+                            <execution>
+                                <id>lancedb-jni</id>
+                                <goals>
+                                    <goal>build</goal>
+                                </goals>
+                                <configuration>
+                                    <path>lancedb-jni</path>
+                                    <!--<release>true</release>-->
+                                    <!-- Copy native libraries to target/classes for runtime access -->
+                                    <copyTo>${project.build.directory}/classes/nativelib</copyTo>
+                                    <copyWithPlatformDir>true</copyWithPlatformDir>
+                                </configuration>
+                            </execution>
+                            <execution>
+                                <id>lancedb-jni-test</id>
+                                <goals>
+                                    <goal>test</goal>
+                                </goals>
+                                <configuration>
+                                    <path>lancedb-jni</path>
+                                    <release>false</release>
+                                    <verbosity>-v</verbosity>
+                                </configuration>
+                            </execution>
+                        </executions>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
+    </profiles>
+</project>

java/core/src/main/java/com/lancedb/lancedb/Connection.java

@@ -0,0 +1,120 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.lancedb.lancedb;
+
+import io.questdb.jar.jni.JarJniLoader;
+import java.io.Closeable;
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * Represents LanceDB database.
+ */
+public class Connection implements Closeable {
+  static {
+    JarJniLoader.loadLib(Connection.class, "/nativelib", "lancedb_jni");
+  }
+
+  private long nativeConnectionHandle;
+
+  /**
+   * Connect to a LanceDB instance.
+   */
+  public static native Connection connect(String uri);
+
+  /**
+   * Get the names of all tables in the database. The names are sorted in
+   * ascending order.
+   *
+   * @return the table names
+   */
+  public List<String> tableNames() {
+    return tableNames(Optional.empty(), Optional.empty());
+  }
+
+  /**
+   * Get the names of filtered tables in the database. The names are sorted in
+   * ascending order.
+   *
+   * @param limit The number of results to return.
+   * @return the table names
+   */
+  public List<String> tableNames(int limit) {
+    return tableNames(Optional.empty(), Optional.of(limit));
+  }
+
+  /**
+   * Get the names of filtered tables in the database. The names are sorted in
+   * ascending order.
+   *
+   * @param startAfter If present, only return names that come lexicographically after the supplied
+   *                   value. This can be combined with limit to implement pagination
+   *                   by setting this to the last table name from the previous page.
+   * @return the table names
+   */
+  public List<String> tableNames(String startAfter) {
+    return tableNames(Optional.of(startAfter), Optional.empty());
+  }
+
+  /**
+   * Get the names of filtered tables in the database. The names are sorted in
+   * ascending order.
+   *
+   * @param startAfter If present, only return names that come lexicographically after the supplied
+   *                   value. This can be combined with limit to implement pagination
+   *                   by setting this to the last table name from the previous page.
+   * @param limit The number of results to return.
+   * @return the table names
+   */
+  public List<String> tableNames(String startAfter, int limit) {
+    return tableNames(Optional.of(startAfter), Optional.of(limit));
+  }
+
+  /**
+   * Get the names of filtered tables in the database. The names are sorted in
+   * ascending order.
+   *
+   * @param startAfter If present, only return names that come lexicographically after the supplied
+   *                   value. This can be combined with limit to implement pagination
+   *                   by setting this to the last table name from the previous page.
+   * @param limit The number of results to return.
+   * @return the table names
+   */
+  public native List<String> tableNames(
+      Optional<String> startAfter, Optional<Integer> limit);
+
+  /**
+   * Closes this connection and releases any system resources associated with it. If
+   * the connection is
+   * already closed, then invoking this method has no effect.
+   */
+  @Override
+  public void close() {
+    if (nativeConnectionHandle != 0) {
+      releaseNativeConnection(nativeConnectionHandle);
+      nativeConnectionHandle = 0;
+    }
+  }
+
+  /**
+   * Native method to release the Lance connection resources associated with the
+   * given handle.
+   *
+   * @param handle The native handle to the connection resource.
+   */
+  private native void releaseNativeConnection(long handle);
+
+  private Connection() {}
+}

java/core/src/test/java/com/lancedb/lancedb/ConnectionTest.java

@@ -0,0 +1,135 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.lancedb.lancedb;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Path;
+import java.util.List;
+import java.net.URL;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+public class ConnectionTest {
+  private static final String[] TABLE_NAMES = {
+      "dataset_version",
+      "new_empty_dataset",
+      "test",
+      "write_stream"
+  };
+
+  @TempDir
+  static Path tempDir; // Temporary directory for the tests
+  private static URL lanceDbURL;
+
+  @BeforeAll
+  static void setUp() {
+    ClassLoader classLoader = ConnectionTest.class.getClassLoader();
+    lanceDbURL = classLoader.getResource("example_db");
+  }
+
+  @Test
+  void emptyDB() {
+    String databaseUri = tempDir.resolve("emptyDB").toString();
+    try (Connection conn = Connection.connect(databaseUri)) {
+      List<String> tableNames = conn.tableNames();
+      assertTrue(tableNames.isEmpty());
+    }
+  }
+
+  @Test
+  void tableNames() {
+    try (Connection conn = Connection.connect(lanceDbURL.toString())) {
+      List<String> tableNames = conn.tableNames();
+      assertEquals(4, tableNames.size());
+      for (int i = 0; i < TABLE_NAMES.length; i++) {
+        assertEquals(TABLE_NAMES[i], tableNames.get(i));
+      }
+    }
+  }
+
+  @Test
+  void tableNamesStartAfter() {
+    try (Connection conn = Connection.connect(lanceDbURL.toString())) {
+      assertTableNamesStartAfter(conn, TABLE_NAMES[0], 3, TABLE_NAMES[1], TABLE_NAMES[2], TABLE_NAMES[3]);
+      assertTableNamesStartAfter(conn, TABLE_NAMES[1], 2, TABLE_NAMES[2], TABLE_NAMES[3]);
+      assertTableNamesStartAfter(conn, TABLE_NAMES[2], 1, TABLE_NAMES[3]);
+      assertTableNamesStartAfter(conn, TABLE_NAMES[3], 0);
+      assertTableNamesStartAfter(conn, "a_dataset", 4, TABLE_NAMES[0], TABLE_NAMES[1], TABLE_NAMES[2], TABLE_NAMES[3]);
+      assertTableNamesStartAfter(conn, "o_dataset", 2, TABLE_NAMES[2], TABLE_NAMES[3]);
+      assertTableNamesStartAfter(conn, "v_dataset", 1, TABLE_NAMES[3]);
+      assertTableNamesStartAfter(conn, "z_dataset", 0);
+    }
+  }
+
+  private void assertTableNamesStartAfter(Connection conn, String startAfter, int expectedSize, String... expectedNames) {
+    List<String> tableNames = conn.tableNames(startAfter);
+    assertEquals(expectedSize, tableNames.size());
+    for (int i = 0; i < expectedNames.length; i++) {
+      assertEquals(expectedNames[i], tableNames.get(i));
+    }
+  }
+
+  @Test
+  void tableNamesLimit() {
+      try (Connection conn = Connection.connect(lanceDbURL.toString())) {
+      for (int i = 0; i <= TABLE_NAMES.length; i++) {
+        List<String> tableNames = conn.tableNames(i);
+        assertEquals(i, tableNames.size());
+        for (int j = 0; j < i; j++) {
+          assertEquals(TABLE_NAMES[j], tableNames.get(j));
+        }
+      }
+    }
+  }
+
+  @Test
+  void tableNamesStartAfterLimit() {
+    try (Connection conn = Connection.connect(lanceDbURL.toString())) {
+      List<String> tableNames = conn.tableNames(TABLE_NAMES[0], 2);
+      assertEquals(2, tableNames.size());
+      assertEquals(TABLE_NAMES[1], tableNames.get(0));
+      assertEquals(TABLE_NAMES[2], tableNames.get(1));
+      tableNames = conn.tableNames(TABLE_NAMES[1], 1);
+      assertEquals(1, tableNames.size());
+      assertEquals(TABLE_NAMES[2], tableNames.get(0));
+      tableNames = conn.tableNames(TABLE_NAMES[2], 2);
+      assertEquals(1, tableNames.size());
+      assertEquals(TABLE_NAMES[3], tableNames.get(0));
+      tableNames = conn.tableNames(TABLE_NAMES[3], 2);
+      assertEquals(0, tableNames.size());
+      tableNames = conn.tableNames(TABLE_NAMES[0], 0);
+      assertEquals(0, tableNames.size());
+
+      // Limit larger than the number of remaining tables
+      tableNames = conn.tableNames(TABLE_NAMES[0], 10);
+      assertEquals(3, tableNames.size());
+      assertEquals(TABLE_NAMES[1], tableNames.get(0));
+      assertEquals(TABLE_NAMES[2], tableNames.get(1));
+      assertEquals(TABLE_NAMES[3], tableNames.get(2));
+
+      // Start after a value not in the list
+      tableNames = conn.tableNames("non_existent_table", 2);
+      assertEquals(2, tableNames.size());
+      assertEquals(TABLE_NAMES[2], tableNames.get(0));
+      assertEquals(TABLE_NAMES[3], tableNames.get(1));
+
+      // Start after the last table with a limit
+      tableNames = conn.tableNames(TABLE_NAMES[3], 1);
+      assertEquals(0, tableNames.size());
+    }
+  }
+}

java/core/src/test/resources/example_db/dataset_version.lance/_transactions/0-d51afd07-e3cd-4c76-9b9b-787e13fd55b0.txn

@@ -0,0 +1 @@
+$d51afd07-e3cd-4c76-9b9b-787e13fd55b0<62>=id <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
*int320
8
name
 <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
*string0
8

java/core/src/test/resources/example_db/new_empty_dataset.lance/_transactions/0-15648e72-076f-4ef1-8b90-10d305b95b3b.txn

@@ -0,0 +1 @@
+$15648e72-076f-4ef1-8b90-10d305b95b3b<33>=id <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
*int320
8
name
 <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
*string0
8

java/core/src/test/resources/example_db/test.lance/_transactions/0-a3689caf-4f6b-4afc-a3c7-97af75661843.txn

@@ -0,0 +1 @@
+$a3689caf-4f6b-4afc-a3c7-97af75661843<34>oitem <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
*string8price
 <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
*double8
0vector <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
*fixed_size_list:float:28

java/pom.xml

@@ -0,0 +1,129 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <groupId>com.lancedb</groupId>
+    <artifactId>lancedb-parent</artifactId>
+    <version>0.1-SNAPSHOT</version>
+    <packaging>pom</packaging>
+
+    <name>Lance Parent</name>
+
+    <properties>
+        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+        <maven.compiler.source>11</maven.compiler.source>
+        <maven.compiler.target>11</maven.compiler.target>
+        <arrow.version>15.0.0</arrow.version>
+    </properties>
+
+    <modules>
+        <module>core</module>
+    </modules>
+
+    <dependencyManagement>
+        <dependencies>
+            <dependency>
+                <groupId>org.apache.arrow</groupId>
+                <artifactId>arrow-vector</artifactId>
+                <version>${arrow.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>org.apache.arrow</groupId>
+                <artifactId>arrow-memory-netty</artifactId>
+                <version>${arrow.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>org.apache.arrow</groupId>
+                <artifactId>arrow-c-data</artifactId>
+                <version>${arrow.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>org.apache.arrow</groupId>
+                <artifactId>arrow-dataset</artifactId>
+                <version>${arrow.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>org.questdb</groupId>
+                <artifactId>jar-jni</artifactId>
+                <version>1.1.1</version>
+            </dependency>
+            <dependency>
+                <groupId>org.junit.jupiter</groupId>
+                <artifactId>junit-jupiter</artifactId>
+                <version>5.10.1</version>
+            </dependency>
+            <dependency>
+                <groupId>org.json</groupId>
+                <artifactId>json</artifactId>
+                <version>20210307</version>
+            </dependency>
+        </dependencies>
+    </dependencyManagement>
+
+     <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-checkstyle-plugin</artifactId>
+                <version>3.3.1</version>
+                <configuration>
+                    <configLocation>google_checks.xml</configLocation>
+                    <consoleOutput>true</consoleOutput>
+                    <failsOnError>true</failsOnError>
+                    <violationSeverity>warning</violationSeverity>
+                    <linkXRef>false</linkXRef>
+                </configuration>
+                <executions>
+                    <execution>
+                        <id>validate</id>
+                        <phase>validate</phase>
+                        <goals>
+                            <goal>check</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+        <pluginManagement>
+            <plugins>
+                <plugin>
+                    <artifactId>maven-clean-plugin</artifactId>
+                    <version>3.1.0</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-resources-plugin</artifactId>
+                    <version>3.0.2</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-compiler-plugin</artifactId>
+                    <version>3.8.1</version>
+                    <configuration>
+                        <compilerArgs>
+                            <arg>-h</arg>
+                            <arg>target/headers</arg>
+                        </compilerArgs>
+                    </configuration>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-surefire-plugin</artifactId>
+                    <version>3.2.5</version>
+                    <configuration>
+                        <argLine>--add-opens=java.base/java.nio=ALL-UNNAMED</argLine>
+                        <forkNode implementation="org.apache.maven.plugin.surefire.extensions.SurefireForkNodeFactory"/>
+                        <useSystemClassLoader>false</useSystemClassLoader>
+                    </configuration>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-jar-plugin</artifactId>
+                    <version>3.0.2</version>
+                </plugin>
+                <plugin>
+                    <artifactId>maven-install-plugin</artifactId>
+                    <version>2.5.2</version>
+                </plugin>
+            </plugins>
+        </pluginManagement>
+    </build>
+</project>

node/package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "vectordb",
-  "version": "0.4.20",
+  "version": "0.5.1",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "vectordb",
-      "version": "0.4.20",
+      "version": "0.5.1",
       "cpu": [
         "x64",
         "arm64"

node/package.json

@@ -1,6 +1,6 @@
 {
   "name": "vectordb",
-  "version": "0.4.20",
+  "version": "0.5.1",
   "description": " Serverless, low-latency vector database for AI applications",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

node/src/arrow.ts

@@ -624,8 +624,6 @@ function validateSchemaEmbeddings(
   }
 
   if (missingEmbeddingFields.length > 0 && embeddings === undefined) {
-    console.log({ missingEmbeddingFields, embeddings });
-
     throw new Error(
       `Table has embeddings: "${missingEmbeddingFields
         .map((f) => f.name)
@@ -633,5 +631,5 @@ function validateSchemaEmbeddings(
     );
   }
 
-  return new Schema(fields);
+  return new Schema(fields, schema.metadata);
 }

node/src/index.ts

@@ -704,6 +704,9 @@ export interface VectorIndex {
 export interface IndexStats {
   numIndexedRows: number | null
   numUnindexedRows: number | null
+  index_type: string | null
+  distance_type: string | null
+  completed_at: string | null
 }
 
 /**

node/src/remote/index.ts

@@ -509,7 +509,8 @@ export class RemoteTable<T = number[]> implements Table<T> {
     return (await results.body()).indexes?.map((index: any) => ({
       columns: index.columns,
       name: index.index_name,
-      uuid: index.index_uuid
+      uuid: index.index_uuid,
+      status: index.status
     }))
   }
 
@@ -520,7 +521,10 @@ export class RemoteTable<T = number[]> implements Table<T> {
     const body = await results.body()
     return {
       numIndexedRows: body?.num_indexed_rows,
-      numUnindexedRows: body?.num_unindexed_rows
+      numUnindexedRows: body?.num_unindexed_rows,
+      index_type: body?.index_type,
+      distance_type: body?.distance_type,
+      completed_at: body?.completed_at
     }
   }
 

nodejs/__test__/arrow.test.ts

@@ -31,6 +31,7 @@ import {
   Schema,
   Struct,
   type Table,
+  Type,
   Utf8,
   tableFromIPC,
 } from "apache-arrow";
@@ -51,7 +52,12 @@ import {
   makeArrowTable,
   makeEmptyTable,
 } from "../lancedb/arrow";
-import { type EmbeddingFunction } from "../lancedb/embedding/embedding_function";
+import {
+  EmbeddingFunction,
+  FieldOptions,
+  FunctionOptions,
+} from "../lancedb/embedding/embedding_function";
+import { EmbeddingFunctionConfig } from "../lancedb/embedding/registry";
 
 // biome-ignore lint/suspicious/noExplicitAny: skip
 function sampleRecords(): Array<Record<string, any>> {
@@ -280,23 +286,46 @@ describe("The function makeArrowTable", function () {
   });
 });
 
-class DummyEmbedding implements EmbeddingFunction<string> {
-  public readonly sourceColumn = "string";
-  public readonly embeddingDimension = 2;
-  public readonly embeddingDataType = new Float16();
+class DummyEmbedding extends EmbeddingFunction<string> {
+  toJSON(): Partial<FunctionOptions> {
+    return {};
+  }
 
-  async embed(data: string[]): Promise<number[][]> {
+  async computeSourceEmbeddings(data: string[]): Promise<number[][]> {
     return data.map(() => [0.0, 0.0]);
   }
+
+  ndims(): number {
+    return 2;
+  }
+
+  embeddingDataType() {
+    return new Float16();
+  }
 }
 
-class DummyEmbeddingWithNoDimension implements EmbeddingFunction<string> {
-  public readonly sourceColumn = "string";
+class DummyEmbeddingWithNoDimension extends EmbeddingFunction<string> {
+  toJSON(): Partial<FunctionOptions> {
+    return {};
+  }
 
-  async embed(data: string[]): Promise<number[][]> {
+  embeddingDataType(): Float {
+    return new Float16();
+  }
+
+  async computeSourceEmbeddings(data: string[]): Promise<number[][]> {
     return data.map(() => [0.0, 0.0]);
   }
 }
+const dummyEmbeddingConfig: EmbeddingFunctionConfig = {
+  sourceColumn: "string",
+  function: new DummyEmbedding(),
+};
+
+const dummyEmbeddingConfigWithNoDimension: EmbeddingFunctionConfig = {
+  sourceColumn: "string",
+  function: new DummyEmbeddingWithNoDimension(),
+};
 
 describe("convertToTable", function () {
   it("will infer data types correctly", async function () {
@@ -331,7 +360,7 @@ describe("convertToTable", function () {
 
   it("will apply embeddings", async function () {
     const records = sampleRecords();
-    const table = await convertToTable(records, new DummyEmbedding());
+    const table = await convertToTable(records, dummyEmbeddingConfig);
     expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe(true);
     expect(table.getChild("vector")?.type.children[0].type.toString()).toEqual(
       new Float16().toString(),
@@ -340,7 +369,7 @@ describe("convertToTable", function () {
 
   it("will fail if missing the embedding source column", async function () {
     await expect(
-      convertToTable([{ id: 1 }], new DummyEmbedding()),
+      convertToTable([{ id: 1 }], dummyEmbeddingConfig),
     ).rejects.toThrow("'string' was not present");
   });
 
@@ -351,7 +380,7 @@ describe("convertToTable", function () {
     const table = makeEmptyTable(schema);
 
     // If the embedding specifies the dimension we are fine
-    await fromTableToBuffer(table, new DummyEmbedding());
+    await fromTableToBuffer(table, dummyEmbeddingConfig);
 
     // We can also supply a schema and should be ok
     const schemaWithEmbedding = new Schema([
@@ -364,13 +393,13 @@ describe("convertToTable", function () {
     ]);
     await fromTableToBuffer(
       table,
-      new DummyEmbeddingWithNoDimension(),
+      dummyEmbeddingConfigWithNoDimension,
       schemaWithEmbedding,
     );
 
     // Otherwise we will get an error
     await expect(
-      fromTableToBuffer(table, new DummyEmbeddingWithNoDimension()),
+      fromTableToBuffer(table, dummyEmbeddingConfigWithNoDimension),
     ).rejects.toThrow("does not specify `embeddingDimension`");
   });
 
@@ -383,7 +412,7 @@ describe("convertToTable", function () {
         false,
       ),
     ]);
-    const table = await convertToTable([], new DummyEmbedding(), { schema });
+    const table = await convertToTable([], dummyEmbeddingConfig, { schema });
     expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe(true);
     expect(table.getChild("vector")?.type.children[0].type.toString()).toEqual(
       new Float16().toString(),
@@ -393,16 +422,17 @@ describe("convertToTable", function () {
   it("will complain if embeddings present but schema missing embedding column", async function () {
     const schema = new Schema([new Field("string", new Utf8(), false)]);
     await expect(
-      convertToTable([], new DummyEmbedding(), { schema }),
+      convertToTable([], dummyEmbeddingConfig, { schema }),
     ).rejects.toThrow("column vector was missing");
   });
 
   it("will provide a nice error if run twice", async function () {
     const records = sampleRecords();
-    const table = await convertToTable(records, new DummyEmbedding());
+    const table = await convertToTable(records, dummyEmbeddingConfig);
+
     // fromTableToBuffer will try and apply the embeddings again
     await expect(
-      fromTableToBuffer(table, new DummyEmbedding()),
+      fromTableToBuffer(table, dummyEmbeddingConfig),
     ).rejects.toThrow("already existed");
   });
 });

nodejs/__test__/connection.test.ts

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import { Field, Float64, Schema } from "apache-arrow";
 import * as tmp from "tmp";
-
-import { Connection, connect } from "../lancedb";
+import { Connection, Table, connect } from "../lancedb";
 
 describe("when connecting", () => {
   let tmpDir: tmp.DirResult;
@@ -87,4 +87,39 @@ describe("given a connection", () => {
     tables = await db.tableNames({ startAfter: "a" });
     expect(tables).toEqual(["b", "c"]);
   });
+
+  it("should create tables in v2 mode", async () => {
+    const db = await connect(tmpDir.name);
+    const data = [...Array(10000).keys()].map((i) => ({ id: i }));
+
+    // Create in v1 mode
+    let table = await db.createTable("test", data);
+
+    const isV2 = async (table: Table) => {
+      const data = await table.query().toArrow({ maxBatchLength: 100000 });
+      console.log(data.batches.length);
+      return data.batches.length < 5;
+    };
+
+    await expect(isV2(table)).resolves.toBe(false);
+
+    // Create in v2 mode
+    table = await db.createTable("test_v2", data, { useLegacyFormat: false });
+
+    await expect(isV2(table)).resolves.toBe(true);
+
+    await table.add(data);
+
+    await expect(isV2(table)).resolves.toBe(true);
+
+    // Create empty in v2 mode
+    const schema = new Schema([new Field("id", new Float64(), true)]);
+
+    table = await db.createEmptyTable("test_v2_empty", schema, {
+      useLegacyFormat: false,
+    });
+
+    await table.add(data);
+    await expect(isV2(table)).resolves.toBe(true);
+  });
 });

nodejs/__test__/embedding.test.ts

@@ -0,0 +1,314 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import * as tmp from "tmp";
+
+import { connect } from "../lancedb";
+import {
+  Field,
+  FixedSizeList,
+  Float,
+  Float16,
+  Float32,
+  Float64,
+  Schema,
+  Utf8,
+} from "../lancedb/arrow";
+import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
+import { getRegistry, register } from "../lancedb/embedding/registry";
+
+describe("embedding functions", () => {
+  let tmpDir: tmp.DirResult;
+  beforeEach(() => {
+    tmpDir = tmp.dirSync({ unsafeCleanup: true });
+  });
+  afterEach(() => {
+    tmpDir.removeCallback();
+    getRegistry().reset();
+  });
+
+  it("should be able to create a table with an embedding function", async () => {
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      toJSON(): object {
+        return {};
+      }
+      ndims() {
+        return 3;
+      }
+      embeddingDataType(): Float {
+        return new Float32();
+      }
+      async computeQueryEmbeddings(_data: string) {
+        return [1, 2, 3];
+      }
+      async computeSourceEmbeddings(data: string[]) {
+        return Array.from({ length: data.length }).fill([
+          1, 2, 3,
+        ]) as number[][];
+      }
+    }
+    const func = new MockEmbeddingFunction();
+    const db = await connect(tmpDir.name);
+    const table = await db.createTable(
+      "test",
+      [
+        { id: 1, text: "hello" },
+        { id: 2, text: "world" },
+      ],
+      {
+        embeddingFunction: {
+          function: func,
+          sourceColumn: "text",
+        },
+      },
+    );
+    // biome-ignore lint/suspicious/noExplicitAny: test
+    const arr = (await table.query().toArray()) as any;
+    expect(arr[0].vector).toBeDefined();
+
+    // we round trip through JSON to make sure the vector properly gets converted to an array
+    // otherwise it'll be a TypedArray or Vector
+    const vector0 = JSON.parse(JSON.stringify(arr[0].vector));
+    expect(vector0).toEqual([1, 2, 3]);
+  });
+
+  it("should be able to create an empty table with an embedding function", async () => {
+    @register()
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      toJSON(): object {
+        return {};
+      }
+      ndims() {
+        return 3;
+      }
+      embeddingDataType(): Float {
+        return new Float32();
+      }
+      async computeQueryEmbeddings(_data: string) {
+        return [1, 2, 3];
+      }
+      async computeSourceEmbeddings(data: string[]) {
+        return Array.from({ length: data.length }).fill([
+          1, 2, 3,
+        ]) as number[][];
+      }
+    }
+    const schema = new Schema([
+      new Field("text", new Utf8(), true),
+      new Field(
+        "vector",
+        new FixedSizeList(3, new Field("item", new Float32(), true)),
+        true,
+      ),
+    ]);
+
+    const func = new MockEmbeddingFunction();
+    const db = await connect(tmpDir.name);
+    const table = await db.createEmptyTable("test", schema, {
+      embeddingFunction: {
+        function: func,
+        sourceColumn: "text",
+      },
+    });
+    const outSchema = await table.schema();
+    expect(outSchema.metadata.get("embedding_functions")).toBeDefined();
+    await table.add([{ text: "hello world" }]);
+
+    // biome-ignore lint/suspicious/noExplicitAny: test
+    const arr = (await table.query().toArray()) as any;
+    expect(arr[0].vector).toBeDefined();
+
+    // we round trip through JSON to make sure the vector properly gets converted to an array
+    // otherwise it'll be a TypedArray or Vector
+    const vector0 = JSON.parse(JSON.stringify(arr[0].vector));
+    expect(vector0).toEqual([1, 2, 3]);
+  });
+  it("should error when appending to a table with an unregistered embedding function", async () => {
+    @register("mock")
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      toJSON(): object {
+        return {};
+      }
+      ndims() {
+        return 3;
+      }
+      embeddingDataType(): Float {
+        return new Float32();
+      }
+      async computeQueryEmbeddings(_data: string) {
+        return [1, 2, 3];
+      }
+      async computeSourceEmbeddings(data: string[]) {
+        return Array.from({ length: data.length }).fill([
+          1, 2, 3,
+        ]) as number[][];
+      }
+    }
+    const func = getRegistry().get<MockEmbeddingFunction>("mock")!.create();
+
+    const schema = LanceSchema({
+      id: new Float64(),
+      text: func.sourceField(new Utf8()),
+      vector: func.vectorField(),
+    });
+
+    const db = await connect(tmpDir.name);
+    await db.createTable(
+      "test",
+      [
+        { id: 1, text: "hello" },
+        { id: 2, text: "world" },
+      ],
+      {
+        schema,
+      },
+    );
+
+    getRegistry().reset();
+    const db2 = await connect(tmpDir.name);
+
+    const tbl = await db2.openTable("test");
+
+    expect(tbl.add([{ id: 3, text: "hello" }])).rejects.toThrow(
+      `Function "mock" not found in registry`,
+    );
+  });
+  test.each([new Float16(), new Float32(), new Float64()])(
+    "should be able to provide manual embeddings with multiple float datatype",
+    async (floatType) => {
+      class MockEmbeddingFunction extends EmbeddingFunction<string> {
+        toJSON(): object {
+          return {};
+        }
+        ndims() {
+          return 3;
+        }
+        embeddingDataType(): Float {
+          return floatType;
+        }
+        async computeQueryEmbeddings(_data: string) {
+          return [1, 2, 3];
+        }
+        async computeSourceEmbeddings(data: string[]) {
+          return Array.from({ length: data.length }).fill([
+            1, 2, 3,
+          ]) as number[][];
+        }
+      }
+      const data = [{ text: "hello" }, { text: "hello world" }];
+
+      const schema = new Schema([
+        new Field("vector", new FixedSizeList(3, new Field("item", floatType))),
+        new Field("text", new Utf8()),
+      ]);
+      const func = new MockEmbeddingFunction();
+
+      const name = "test";
+      const db = await connect(tmpDir.name);
+
+      const table = await db.createTable(name, data, {
+        schema,
+        embeddingFunction: {
+          sourceColumn: "text",
+          function: func,
+        },
+      });
+      const res = await table.query().toArray();
+
+      expect([...res[0].vector]).toEqual([1, 2, 3]);
+    },
+  );
+
+  test.only.each([new Float16(), new Float32(), new Float64()])(
+    "should be able to provide auto embeddings with multiple float datatypes",
+    async (floatType) => {
+      @register("test1")
+      class MockEmbeddingFunctionWithoutNDims extends EmbeddingFunction<string> {
+        toJSON(): object {
+          return {};
+        }
+
+        embeddingDataType(): Float {
+          return floatType;
+        }
+        async computeQueryEmbeddings(_data: string) {
+          return [1, 2, 3];
+        }
+        async computeSourceEmbeddings(data: string[]) {
+          return Array.from({ length: data.length }).fill([
+            1, 2, 3,
+          ]) as number[][];
+        }
+      }
+      @register("test")
+      class MockEmbeddingFunction extends EmbeddingFunction<string> {
+        toJSON(): object {
+          return {};
+        }
+        ndims() {
+          return 3;
+        }
+        embeddingDataType(): Float {
+          return floatType;
+        }
+        async computeQueryEmbeddings(_data: string) {
+          return [1, 2, 3];
+        }
+        async computeSourceEmbeddings(data: string[]) {
+          return Array.from({ length: data.length }).fill([
+            1, 2, 3,
+          ]) as number[][];
+        }
+      }
+      const func = getRegistry().get<MockEmbeddingFunction>("test")!.create();
+      const func2 = getRegistry()
+        .get<MockEmbeddingFunctionWithoutNDims>("test1")!
+        .create();
+
+      const schema = LanceSchema({
+        text: func.sourceField(new Utf8()),
+        vector: func.vectorField(floatType),
+      });
+
+      const schema2 = LanceSchema({
+        text: func2.sourceField(new Utf8()),
+        vector: func2.vectorField({ datatype: floatType, dims: 3 }),
+      });
+      const schema3 = LanceSchema({
+        text: func2.sourceField(new Utf8()),
+        vector: func.vectorField({
+          datatype: new FixedSizeList(3, new Field("item", floatType, true)),
+          dims: 3,
+        }),
+      });
+
+      const expectedSchema = new Schema([
+        new Field("text", new Utf8(), true),
+        new Field(
+          "vector",
+          new FixedSizeList(3, new Field("item", floatType, true)),
+          true,
+        ),
+      ]);
+      const stringSchema = JSON.stringify(schema, null, 2);
+      const stringSchema2 = JSON.stringify(schema2, null, 2);
+      const stringSchema3 = JSON.stringify(schema3, null, 2);
+      const stringExpectedSchema = JSON.stringify(expectedSchema, null, 2);
+
+      expect(stringSchema).toEqual(stringExpectedSchema);
+      expect(stringSchema2).toEqual(stringExpectedSchema);
+      expect(stringSchema3).toEqual(stringExpectedSchema);
+    },
+  );
+});

nodejs/__test__/registry.test.ts

@@ -0,0 +1,169 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+import * as arrow from "apache-arrow";
+import * as arrowOld from "apache-arrow-old";
+
+import * as tmp from "tmp";
+
+import { connect } from "../lancedb";
+import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
+import { getRegistry, register } from "../lancedb/embedding/registry";
+
+describe.each([arrow, arrowOld])("LanceSchema", (arrow) => {
+  test("should preserve input order", async () => {
+    const schema = LanceSchema({
+      id: new arrow.Int32(),
+      text: new arrow.Utf8(),
+      vector: new arrow.Float32(),
+    });
+    expect(schema.fields.map((x) => x.name)).toEqual(["id", "text", "vector"]);
+  });
+});
+
+describe("Registry", () => {
+  let tmpDir: tmp.DirResult;
+  beforeEach(() => {
+    tmpDir = tmp.dirSync({ unsafeCleanup: true });
+  });
+
+  afterEach(() => {
+    tmpDir.removeCallback();
+    getRegistry().reset();
+  });
+
+  it("should register a new item to the registry", async () => {
+    @register("mock-embedding")
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      toJSON(): object {
+        return {
+          someText: "hello",
+        };
+      }
+      constructor() {
+        super();
+      }
+      ndims() {
+        return 3;
+      }
+      embeddingDataType(): arrow.Float {
+        return new arrow.Float32();
+      }
+      async computeSourceEmbeddings(data: string[]) {
+        return data.map(() => [1, 2, 3]);
+      }
+    }
+    const func = getRegistry()
+      .get<MockEmbeddingFunction>("mock-embedding")!
+      .create();
+
+    const schema = LanceSchema({
+      id: new arrow.Int32(),
+      text: func.sourceField(new arrow.Utf8()),
+      vector: func.vectorField(),
+    });
+
+    const db = await connect(tmpDir.name);
+    const table = await db.createTable(
+      "test",
+      [
+        { id: 1, text: "hello" },
+        { id: 2, text: "world" },
+      ],
+      { schema },
+    );
+    const expected = [
+      [1, 2, 3],
+      [1, 2, 3],
+    ];
+    const actual = await table.query().toArrow();
+    const vectors = actual
+      .getChild("vector")
+      ?.toArray()
+      .map((x: unknown) => {
+        if (x instanceof arrow.Vector) {
+          return [...x];
+        } else {
+          return x;
+        }
+      });
+    expect(vectors).toEqual(expected);
+  });
+  test("should error if registering with the same name", async () => {
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      toJSON(): object {
+        return {
+          someText: "hello",
+        };
+      }
+      constructor() {
+        super();
+      }
+      ndims() {
+        return 3;
+      }
+      embeddingDataType(): arrow.Float {
+        return new arrow.Float32();
+      }
+      async computeSourceEmbeddings(data: string[]) {
+        return data.map(() => [1, 2, 3]);
+      }
+    }
+    register("mock-embedding")(MockEmbeddingFunction);
+    expect(() => register("mock-embedding")(MockEmbeddingFunction)).toThrow(
+      'Embedding function with alias "mock-embedding" already exists',
+    );
+  });
+  test("schema should contain correct metadata", async () => {
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      toJSON(): object {
+        return {
+          someText: "hello",
+        };
+      }
+      constructor() {
+        super();
+      }
+      ndims() {
+        return 3;
+      }
+      embeddingDataType(): arrow.Float {
+        return new arrow.Float32();
+      }
+      async computeSourceEmbeddings(data: string[]) {
+        return data.map(() => [1, 2, 3]);
+      }
+    }
+    const func = new MockEmbeddingFunction();
+
+    const schema = LanceSchema({
+      id: new arrow.Int32(),
+      text: func.sourceField(new arrow.Utf8()),
+      vector: func.vectorField(),
+    });
+    const expectedMetadata = new Map<string, string>([
+      [
+        "embedding_functions",
+        JSON.stringify([
+          {
+            sourceColumn: "text",
+            vectorColumn: "vector",
+            name: "MockEmbeddingFunction",
+            model: { someText: "hello" },
+          },
+        ]),
+      ],
+    ]);
+    expect(schema.metadata).toEqual(expectedMetadata);
+  });
+});

nodejs/__test__/table.test.ts

@@ -16,7 +16,12 @@ import * as fs from "fs";
 import * as path from "path";
 import * as tmp from "tmp";
 
+import * as arrow from "apache-arrow";
+import * as arrowOld from "apache-arrow-old";
+
+import { Table, connect } from "../lancedb";
 import {
+  Table as ArrowTable,
   Field,
   FixedSizeList,
   Float32,
@@ -24,15 +29,20 @@ import {
   Int32,
   Int64,
   Schema,
-} from "apache-arrow";
-import { Table, connect } from "../lancedb";
-import { makeArrowTable } from "../lancedb/arrow";
+  makeArrowTable,
+} from "../lancedb/arrow";
+import { EmbeddingFunction, LanceSchema, register } from "../lancedb/embedding";
 import { Index } from "../lancedb/indices";
 
-describe("Given a table", () => {
+// biome-ignore lint/suspicious/noExplicitAny: <explanation>
+describe.each([arrow, arrowOld])("Given a table", (arrow: any) => {
   let tmpDir: tmp.DirResult;
   let table: Table;
-  const schema = new Schema([new Field("id", new Float64(), true)]);
+
+  const schema = new arrow.Schema([
+    new arrow.Field("id", new arrow.Float64(), true),
+  ]);
+
   beforeEach(async () => {
     tmpDir = tmp.dirSync({ unsafeCleanup: true });
     const conn = await connect(tmpDir.name);
@@ -83,6 +93,43 @@ describe("Given a table", () => {
     expect(await table.countRows("id == 7")).toBe(1);
     expect(await table.countRows("id == 10")).toBe(1);
   });
+
+  // https://github.com/lancedb/lancedb/issues/1293
+  test.each([new arrow.Float16(), new arrow.Float32(), new arrow.Float64()])(
+    "can create empty table with non default float type: %s",
+    async (floatType) => {
+      const db = await connect(tmpDir.name);
+
+      const data = [
+        { text: "hello", vector: Array(512).fill(1.0) },
+        { text: "hello world", vector: Array(512).fill(1.0) },
+      ];
+      const f64Schema = new arrow.Schema([
+        new arrow.Field("text", new arrow.Utf8(), true),
+        new arrow.Field(
+          "vector",
+          new arrow.FixedSizeList(512, new arrow.Field("item", floatType)),
+          true,
+        ),
+      ]);
+
+      const f64Table = await db.createEmptyTable("f64", f64Schema, {
+        mode: "overwrite",
+      });
+      try {
+        await f64Table.add(data);
+        const res = await f64Table.query().toArray();
+        expect(res.length).toBe(2);
+      } catch (e) {
+        expect(e).toBeUndefined();
+      }
+    },
+  );
+
+  it("should return the table as an instance of an arrow table", async () => {
+    const arrowTbl = await table.toArrow();
+    expect(arrowTbl).toBeInstanceOf(ArrowTable);
+  });
 });
 
 describe("When creating an index", () => {
@@ -419,3 +466,127 @@ describe("when dealing with versioning", () => {
     );
   });
 });
+
+describe("when optimizing a dataset", () => {
+  let tmpDir: tmp.DirResult;
+  let table: Table;
+  beforeEach(async () => {
+    tmpDir = tmp.dirSync({ unsafeCleanup: true });
+    const con = await connect(tmpDir.name);
+    table = await con.createTable("vectors", [{ id: 1 }]);
+    await table.add([{ id: 2 }]);
+  });
+  afterEach(() => {
+    tmpDir.removeCallback();
+  });
+
+  it("compacts files", async () => {
+    const stats = await table.optimize();
+    expect(stats.compaction.filesAdded).toBe(1);
+    expect(stats.compaction.filesRemoved).toBe(2);
+    expect(stats.compaction.fragmentsAdded).toBe(1);
+    expect(stats.compaction.fragmentsRemoved).toBe(2);
+  });
+
+  it("cleanups old versions", async () => {
+    const stats = await table.optimize({ cleanupOlderThan: new Date() });
+    expect(stats.prune.bytesRemoved).toBeGreaterThan(0);
+    expect(stats.prune.oldVersionsRemoved).toBe(3);
+  });
+});
+
+describe("table.search", () => {
+  let tmpDir: tmp.DirResult;
+  beforeEach(() => {
+    tmpDir = tmp.dirSync({ unsafeCleanup: true });
+  });
+  afterEach(() => tmpDir.removeCallback());
+
+  test("can search using a string", async () => {
+    @register()
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      toJSON(): object {
+        return {};
+      }
+      ndims() {
+        return 1;
+      }
+      embeddingDataType(): arrow.Float {
+        return new Float32();
+      }
+
+      // Hardcoded embeddings for the sake of testing
+      async computeQueryEmbeddings(_data: string) {
+        switch (_data) {
+          case "greetings":
+            return [0.1];
+          case "farewell":
+            return [0.2];
+          default:
+            return null as never;
+        }
+      }
+
+      // Hardcoded embeddings for the sake of testing
+      async computeSourceEmbeddings(data: string[]) {
+        return data.map((s) => {
+          switch (s) {
+            case "hello world":
+              return [0.1];
+            case "goodbye world":
+              return [0.2];
+            default:
+              return null as never;
+          }
+        });
+      }
+    }
+
+    const func = new MockEmbeddingFunction();
+    const schema = LanceSchema({
+      text: func.sourceField(new arrow.Utf8()),
+      vector: func.vectorField(),
+    });
+    const db = await connect(tmpDir.name);
+    const data = [{ text: "hello world" }, { text: "goodbye world" }];
+    const table = await db.createTable("test", data, { schema });
+
+    const results = await table.search("greetings").then((r) => r.toArray());
+    expect(results[0].text).toBe(data[0].text);
+
+    const results2 = await table.search("farewell").then((r) => r.toArray());
+    expect(results2[0].text).toBe(data[1].text);
+  });
+
+  test("rejects if no embedding function provided", async () => {
+    const db = await connect(tmpDir.name);
+    const data = [
+      { text: "hello world", vector: [0.1, 0.2, 0.3] },
+      { text: "goodbye world", vector: [0.4, 0.5, 0.6] },
+    ];
+    const table = await db.createTable("test", data);
+
+    expect(table.search("hello")).rejects.toThrow(
+      "No embedding functions are defined in the table",
+    );
+  });
+
+  test.each([
+    [0.4, 0.5, 0.599], // number[]
+    Float32Array.of(0.4, 0.5, 0.599), // Float32Array
+    Float64Array.of(0.4, 0.5, 0.599), // Float64Array
+  ])("can search using vectorlike datatypes", async (vectorlike) => {
+    const db = await connect(tmpDir.name);
+    const data = [
+      { text: "hello world", vector: [0.1, 0.2, 0.3] },
+      { text: "goodbye world", vector: [0.4, 0.5, 0.6] },
+    ];
+    const table = await db.createTable("test", data);
+
+    // biome-ignore lint/suspicious/noExplicitAny: test
+    const results: any[] = await table.search(vectorlike).toArray();
+
+    expect(results.length).toBe(2);
+    expect(results[0].text).toBe(data[1].text);
+  });
+});

nodejs/biome.json

@@ -48,7 +48,7 @@
         "noUnsafeFinally": "error",
         "noUnsafeOptionalChaining": "error",
         "noUnusedLabels": "error",
-        "noUnusedVariables": "error",
+        "noUnusedVariables": "warn",
         "useIsNan": "error",
         "useValidForDirection": "error",
         "useYield": "error"
@@ -101,7 +101,13 @@
   },
   "overrides": [
     {
-      "include": ["**/*.ts", "**/*.tsx", "**/*.mts", "**/*.cts"],
+      "include": [
+        "**/*.ts",
+        "**/*.tsx",
+        "**/*.mts",
+        "**/*.cts",
+        "__test__/*.test.ts"
+      ],
       "linter": {
         "rules": {
           "correctness": {

nodejs/lancedb/arrow.ts

@@ -17,24 +17,122 @@ import {
   Binary,
   DataType,
   Field,
+  FixedSizeBinary,
   FixedSizeList,
-  type Float,
+  Float,
   Float32,
+  Int,
+  LargeBinary,
   List,
+  Null,
   RecordBatch,
   RecordBatchFileWriter,
   RecordBatchStreamWriter,
   Schema,
   Struct,
   Utf8,
-  type Vector,
+  Vector,
   makeBuilder,
   makeData,
   type makeTable,
   vectorFromArray,
 } from "apache-arrow";
 import { type EmbeddingFunction } from "./embedding/embedding_function";
-import { sanitizeSchema } from "./sanitize";
+import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
+import { sanitizeField, sanitizeSchema, sanitizeType } from "./sanitize";
+export * from "apache-arrow";
+
+export type IntoVector = Float32Array | Float64Array | number[];
+
+export function isArrowTable(value: object): value is ArrowTable {
+  if (value instanceof ArrowTable) return true;
+  return "schema" in value && "batches" in value;
+}
+
+export function isDataType(value: unknown): value is DataType {
+  return (
+    value instanceof DataType ||
+    DataType.isNull(value) ||
+    DataType.isInt(value) ||
+    DataType.isFloat(value) ||
+    DataType.isBinary(value) ||
+    DataType.isLargeBinary(value) ||
+    DataType.isUtf8(value) ||
+    DataType.isLargeUtf8(value) ||
+    DataType.isBool(value) ||
+    DataType.isDecimal(value) ||
+    DataType.isDate(value) ||
+    DataType.isTime(value) ||
+    DataType.isTimestamp(value) ||
+    DataType.isInterval(value) ||
+    DataType.isDuration(value) ||
+    DataType.isList(value) ||
+    DataType.isStruct(value) ||
+    DataType.isUnion(value) ||
+    DataType.isFixedSizeBinary(value) ||
+    DataType.isFixedSizeList(value) ||
+    DataType.isMap(value) ||
+    DataType.isDictionary(value)
+  );
+}
+export function isNull(value: unknown): value is Null {
+  return value instanceof Null || DataType.isNull(value);
+}
+export function isInt(value: unknown): value is Int {
+  return value instanceof Int || DataType.isInt(value);
+}
+export function isFloat(value: unknown): value is Float {
+  return value instanceof Float || DataType.isFloat(value);
+}
+export function isBinary(value: unknown): value is Binary {
+  return value instanceof Binary || DataType.isBinary(value);
+}
+export function isLargeBinary(value: unknown): value is LargeBinary {
+  return value instanceof LargeBinary || DataType.isLargeBinary(value);
+}
+export function isUtf8(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isUtf8(value);
+}
+export function isLargeUtf8(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isLargeUtf8(value);
+}
+export function isBool(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isBool(value);
+}
+export function isDecimal(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isDecimal(value);
+}
+export function isDate(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isDate(value);
+}
+export function isTime(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isTime(value);
+}
+export function isTimestamp(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isTimestamp(value);
+}
+export function isInterval(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isInterval(value);
+}
+export function isDuration(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isDuration(value);
+}
+export function isList(value: unknown): value is List {
+  return value instanceof List || DataType.isList(value);
+}
+export function isStruct(value: unknown): value is Struct {
+  return value instanceof Struct || DataType.isStruct(value);
+}
+export function isUnion(value: unknown): value is Struct {
+  return value instanceof Struct || DataType.isUnion(value);
+}
+export function isFixedSizeBinary(value: unknown): value is FixedSizeBinary {
+  return value instanceof FixedSizeBinary || DataType.isFixedSizeBinary(value);
+}
+
+export function isFixedSizeList(value: unknown): value is FixedSizeList {
+  return value instanceof FixedSizeList || DataType.isFixedSizeList(value);
+}
 
 /** Data type accepted by NodeJS SDK */
 export type Data = Record<string, unknown>[] | ArrowTable;
@@ -86,6 +184,7 @@ export class MakeArrowTableOptions {
     vector: new VectorColumnOptions(),
   };
   embeddings?: EmbeddingFunction<unknown>;
+  embeddingFunction?: EmbeddingFunctionConfig;
 
   /**
    * If true then string columns will be encoded with dictionary encoding
@@ -198,6 +297,7 @@ export class MakeArrowTableOptions {
 export function makeArrowTable(
   data: Array<Record<string, unknown>>,
   options?: Partial<MakeArrowTableOptions>,
+  metadata?: Map<string, string>,
 ): ArrowTable {
   if (
     data.length === 0 &&
@@ -209,7 +309,11 @@ export function makeArrowTable(
   const opt = new MakeArrowTableOptions(options !== undefined ? options : {});
   if (opt.schema !== undefined && opt.schema !== null) {
     opt.schema = sanitizeSchema(opt.schema);
-    opt.schema = validateSchemaEmbeddings(opt.schema, data, opt.embeddings);
+    opt.schema = validateSchemaEmbeddings(
+      opt.schema,
+      data,
+      options?.embeddingFunction,
+    );
   }
   const columns: Record<string, Vector> = {};
   // TODO: sample dataset to find missing columns
@@ -290,20 +394,41 @@ export function makeArrowTable(
     // `new ArrowTable(schema, batches)` which does not do any schema inference
     const firstTable = new ArrowTable(columns);
     const batchesFixed = firstTable.batches.map(
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       (batch) => new RecordBatch(opt.schema!, batch.data),
     );
-    return new ArrowTable(opt.schema, batchesFixed);
-  } else {
-    return new ArrowTable(columns);
+    let schema: Schema;
+    if (metadata !== undefined) {
+      let schemaMetadata = opt.schema.metadata;
+      if (schemaMetadata.size === 0) {
+        schemaMetadata = metadata;
+      } else {
+        for (const [key, entry] of schemaMetadata.entries()) {
+          schemaMetadata.set(key, entry);
+        }
+      }
+
+      schema = new Schema(opt.schema.fields, schemaMetadata);
+    } else {
+      schema = opt.schema;
+    }
+    return new ArrowTable(schema, batchesFixed);
   }
+  const tbl = new ArrowTable(columns);
+  if (metadata !== undefined) {
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    (<any>tbl.schema).metadata = metadata;
+  }
+  return tbl;
 }
 
 /**
  * Create an empty Arrow table with the provided schema
  */
-export function makeEmptyTable(schema: Schema): ArrowTable {
-  return makeArrowTable([], { schema });
+export function makeEmptyTable(
+  schema: Schema,
+  metadata?: Map<string, string>,
+): ArrowTable {
+  return makeArrowTable([], { schema }, metadata);
 }
 
 /**
@@ -375,13 +500,74 @@ function makeVector(
   }
 }
 
+/** Helper function to apply embeddings from metadata to an input table */
+async function applyEmbeddingsFromMetadata(
+  table: ArrowTable,
+  schema: Schema,
+): Promise<ArrowTable> {
+  const registry = getRegistry();
+  const functions = registry.parseFunctions(schema.metadata);
+
+  const columns = Object.fromEntries(
+    table.schema.fields.map((field) => [
+      field.name,
+      table.getChild(field.name)!,
+    ]),
+  );
+
+  for (const functionEntry of functions.values()) {
+    const sourceColumn = columns[functionEntry.sourceColumn];
+    const destColumn = functionEntry.vectorColumn ?? "vector";
+    if (sourceColumn === undefined) {
+      throw new Error(
+        `Cannot apply embedding function because the source column '${functionEntry.sourceColumn}' was not present in the data`,
+      );
+    }
+    if (columns[destColumn] !== undefined) {
+      throw new Error(
+        `Attempt to apply embeddings to table failed because column ${destColumn} already existed`,
+      );
+    }
+    if (table.batches.length > 1) {
+      throw new Error(
+        "Internal error: `makeArrowTable` unexpectedly created a table with more than one batch",
+      );
+    }
+    const values = sourceColumn.toArray();
+
+    const vectors =
+      await functionEntry.function.computeSourceEmbeddings(values);
+    if (vectors.length !== values.length) {
+      throw new Error(
+        "Embedding function did not return an embedding for each input element",
+      );
+    }
+    let destType: DataType;
+    const dtype = schema.fields.find((f) => f.name === destColumn)!.type;
+    if (isFixedSizeList(dtype)) {
+      destType = sanitizeType(dtype);
+    } else {
+      throw new Error(
+        "Expected FixedSizeList as datatype for vector field, instead got: " +
+          dtype,
+      );
+    }
+    const vector = makeVector(vectors, destType);
+    columns[destColumn] = vector;
+  }
+  const newTable = new ArrowTable(columns);
+  return alignTable(newTable, schema);
+}
+
 /** Helper function to apply embeddings to an input table */
 async function applyEmbeddings<T>(
   table: ArrowTable,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<ArrowTable> {
-  if (embeddings == null) {
+  if (schema?.metadata.has("embedding_functions")) {
+    return applyEmbeddingsFromMetadata(table, schema!);
+  } else if (embeddings == null || embeddings === undefined) {
     return table;
   }
 
@@ -399,8 +585,9 @@ async function applyEmbeddings<T>(
   const newColumns = Object.fromEntries(colEntries);
 
   const sourceColumn = newColumns[embeddings.sourceColumn];
-  const destColumn = embeddings.destColumn ?? "vector";
-  const innerDestType = embeddings.embeddingDataType ?? new Float32();
+  const destColumn = embeddings.vectorColumn ?? "vector";
+  const innerDestType =
+    embeddings.function.embeddingDataType() ?? new Float32();
   if (sourceColumn === undefined) {
     throw new Error(
       `Cannot apply embedding function because the source column '${embeddings.sourceColumn}' was not present in the data`,
@@ -414,11 +601,9 @@ async function applyEmbeddings<T>(
       // if we call convertToTable with 0 records and a schema that includes the embedding
       return table;
     }
-    if (embeddings.embeddingDimension !== undefined) {
-      const destType = newVectorType(
-        embeddings.embeddingDimension,
-        innerDestType,
-      );
+    const dimensions = embeddings.function.ndims();
+    if (dimensions !== undefined) {
+      const destType = newVectorType(dimensions, innerDestType);
       newColumns[destColumn] = makeVector([], destType);
     } else if (schema != null) {
       const destField = schema.fields.find((f) => f.name === destColumn);
@@ -446,7 +631,9 @@ async function applyEmbeddings<T>(
       );
     }
     const values = sourceColumn.toArray();
-    const vectors = await embeddings.embed(values as T[]);
+    const vectors = await embeddings.function.computeSourceEmbeddings(
+      values as T[],
+    );
     if (vectors.length !== values.length) {
       throw new Error(
         "Embedding function did not return an embedding for each input element",
@@ -486,9 +673,9 @@ async function applyEmbeddings<T>(
  * embedding columns.  If no schema is provded then embedding columns will
  * be placed at the end of the table, after all of the input columns.
  */
-export async function convertToTable<T>(
+export async function convertToTable(
   data: Array<Record<string, unknown>>,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   makeTableOptions?: Partial<MakeArrowTableOptions>,
 ): Promise<ArrowTable> {
   const table = makeArrowTable(data, makeTableOptions);
@@ -496,13 +683,13 @@ export async function convertToTable<T>(
 }
 
 /** Creates the Arrow Type for a Vector column with dimension `dim` */
-function newVectorType<T extends Float>(
+export function newVectorType<T extends Float>(
   dim: number,
   innerType: T,
 ): FixedSizeList<T> {
   // in Lance we always default to have the elements nullable, so we need to set it to true
   // otherwise we often get schema mismatches because the stored data always has schema with nullable elements
-  const children = new Field<T>("item", innerType, true);
+  const children = new Field("item", <T>sanitizeType(innerType), true);
   return new FixedSizeList(dim, children);
 }
 
@@ -513,9 +700,9 @@ function newVectorType<T extends Float>(
  *
  * `schema` is required if data is empty
  */
-export async function fromRecordsToBuffer<T>(
+export async function fromRecordsToBuffer(
   data: Array<Record<string, unknown>>,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
@@ -533,9 +720,9 @@ export async function fromRecordsToBuffer<T>(
  *
  * `schema` is required if data is empty
  */
-export async function fromRecordsToStreamBuffer<T>(
+export async function fromRecordsToStreamBuffer(
   data: Array<Record<string, unknown>>,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
@@ -554,9 +741,9 @@ export async function fromRecordsToStreamBuffer<T>(
  *
  * `schema` is required if the table is empty
  */
-export async function fromTableToBuffer<T>(
+export async function fromTableToBuffer(
   table: ArrowTable,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
@@ -575,19 +762,19 @@ export async function fromTableToBuffer<T>(
  *
  * `schema` is required if the table is empty
  */
-export async function fromDataToBuffer<T>(
+export async function fromDataToBuffer(
   data: Data,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
     schema = sanitizeSchema(schema);
   }
-  if (data instanceof ArrowTable) {
+  if (isArrowTable(data)) {
     return fromTableToBuffer(data, embeddings, schema);
   } else {
-    const table = await convertToTable(data);
-    return fromTableToBuffer(table, embeddings, schema);
+    const table = await convertToTable(data, embeddings, { schema });
+    return fromTableToBuffer(table);
   }
 }
 
@@ -599,9 +786,9 @@ export async function fromDataToBuffer<T>(
  *
  * `schema` is required if the table is empty
  */
-export async function fromTableToStreamBuffer<T>(
+export async function fromTableToStreamBuffer(
   table: ArrowTable,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   const tableWithEmbeddings = await applyEmbeddings(table, embeddings, schema);
@@ -654,7 +841,7 @@ export function createEmptyTable(schema: Schema): ArrowTable {
 function validateSchemaEmbeddings(
   schema: Schema,
   data: Array<Record<string, unknown>>,
-  embeddings: EmbeddingFunction<unknown> | undefined,
+  embeddings: EmbeddingFunctionConfig | undefined,
 ) {
   const fields = [];
   const missingEmbeddingFields = [];
@@ -664,10 +851,25 @@ function validateSchemaEmbeddings(
   // if it does not, we add it to the list of missing embedding fields
   // Finally, we check if those missing embedding fields are `this._embeddings`
   // if they are not, we throw an error
-  for (const field of schema.fields) {
-    if (field.type instanceof FixedSizeList) {
+  for (let field of schema.fields) {
+    if (isFixedSizeList(field.type)) {
+      field = sanitizeField(field);
+
       if (data.length !== 0 && data?.[0]?.[field.name] === undefined) {
-        missingEmbeddingFields.push(field);
+        if (schema.metadata.has("embedding_functions")) {
+          const embeddings = JSON.parse(
+            schema.metadata.get("embedding_functions")!,
+          );
+          if (
+            // biome-ignore lint/suspicious/noExplicitAny: we don't know the type of `f`
+            embeddings.find((f: any) => f["vectorColumn"] === field.name) ===
+            undefined
+          ) {
+            missingEmbeddingFields.push(field);
+          }
+        } else {
+          missingEmbeddingFields.push(field);
+        }
       } else {
         fields.push(field);
       }
@@ -677,8 +879,6 @@ function validateSchemaEmbeddings(
   }
 
   if (missingEmbeddingFields.length > 0 && embeddings === undefined) {
-    console.log({ missingEmbeddingFields, embeddings });
-
     throw new Error(
       `Table has embeddings: "${missingEmbeddingFields
         .map((f) => f.name)
@@ -686,5 +886,5 @@ function validateSchemaEmbeddings(
     );
   }
 
-  return new Schema(fields);
+  return new Schema(fields, schema.metadata);
 }

nodejs/lancedb/connection.ts

@@ -12,8 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Table as ArrowTable, Schema } from "apache-arrow";
-import { fromTableToBuffer, makeArrowTable, makeEmptyTable } from "./arrow";
+import { Table as ArrowTable, Schema } from "./arrow";
+import {
+  fromTableToBuffer,
+  isArrowTable,
+  makeArrowTable,
+  makeEmptyTable,
+} from "./arrow";
+import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
 import { ConnectionOptions, Connection as LanceDbConnection } from "./native";
 import { Table } from "./table";
 
@@ -65,6 +71,14 @@ export interface CreateTableOptions {
    * The available options are described at https://lancedb.github.io/lancedb/guides/storage/
    */
   storageOptions?: Record<string, string>;
+  /**
+   * If true then data files will be written with the legacy format
+   *
+   * The default is true while the new format is in beta
+   */
+  useLegacyFormat?: boolean;
+  schema?: Schema;
+  embeddingFunction?: EmbeddingFunctionConfig;
 }
 
 export interface OpenTableOptions {
@@ -174,6 +188,7 @@ export class Connection {
       cleanseStorageOptions(options?.storageOptions),
       options?.indexCacheSize,
     );
+
     return new Table(innerTable);
   }
 
@@ -196,18 +211,25 @@ export class Connection {
     }
 
     let table: ArrowTable;
-    if (data instanceof ArrowTable) {
+    if (isArrowTable(data)) {
       table = data;
     } else {
-      table = makeArrowTable(data);
+      table = makeArrowTable(data, options);
     }
-    const buf = await fromTableToBuffer(table);
+
+    const buf = await fromTableToBuffer(
+      table,
+      options?.embeddingFunction,
+      options?.schema,
+    );
     const innerTable = await this.inner.createTable(
       name,
       buf,
       mode,
       cleanseStorageOptions(options?.storageOptions),
+      options?.useLegacyFormat,
     );
+
     return new Table(innerTable);
   }
 
@@ -227,14 +249,21 @@ export class Connection {
     if (mode === "create" && existOk) {
       mode = "exist_ok";
     }
+    let metadata: Map<string, string> | undefined = undefined;
+    if (options?.embeddingFunction !== undefined) {
+      const embeddingFunction = options.embeddingFunction;
+      const registry = getRegistry();
+      metadata = registry.getTableMetadata([embeddingFunction]);
+    }
 
-    const table = makeEmptyTable(schema);
+    const table = makeEmptyTable(schema, metadata);
     const buf = await fromTableToBuffer(table);
     const innerTable = await this.inner.createEmptyTable(
       name,
       buf,
       mode,
       cleanseStorageOptions(options?.storageOptions),
+      options?.useLegacyFormat,
     );
     return new Table(innerTable);
   }

nodejs/lancedb/embedding/embedding_function.ts

@@ -1,4 +1,4 @@
-// Copyright 2023 Lance Developers.
+// Copyright 2024 Lance Developers.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,67 +12,172 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { type Float } from "apache-arrow";
+import "reflect-metadata";
+import {
+  DataType,
+  Field,
+  FixedSizeList,
+  Float,
+  Float32,
+  type IntoVector,
+  isDataType,
+  isFixedSizeList,
+  isFloat,
+  newVectorType,
+} from "../arrow";
+import { sanitizeType } from "../sanitize";
+
+/**
+ * Options for a given embedding function
+ */
+export interface FunctionOptions {
+  // biome-ignore lint/suspicious/noExplicitAny: options can be anything
+  [key: string]: any;
+}
 
 /**
  * An embedding function that automatically creates vector representation for a given column.
  */
-export interface EmbeddingFunction<T> {
+export abstract class EmbeddingFunction<
+  // biome-ignore lint/suspicious/noExplicitAny: we don't know what the implementor will do
+  T = any,
+  M extends FunctionOptions = FunctionOptions,
+> {
   /**
-   * The name of the column that will be used as input for the Embedding Function.
+   * Convert the embedding function to a JSON object
+   * It is used to serialize the embedding function to the schema
+   * It's important that any object returned by this method contains all the necessary
+   * information to recreate the embedding function
+   *
+   * It should return the same object that was passed to the constructor
+   * If it does not, the embedding function will not be able to be recreated, or could be recreated incorrectly
+   *
+   * @example
+   * ```ts
+   * class MyEmbeddingFunction extends EmbeddingFunction {
+   *   constructor(options: {model: string, timeout: number}) {
+   *     super();
+   *     this.model = options.model;
+   *     this.timeout = options.timeout;
+   *   }
+   *   toJSON() {
+   *     return {
+   *       model: this.model,
+   *       timeout: this.timeout,
+   *     };
+   * }
+   * ```
    */
-  sourceColumn: string;
+  abstract toJSON(): Partial<M>;
 
   /**
-   * The data type of the embedding
+   * sourceField is used in combination with `LanceSchema` to provide a declarative data model
    *
-   * The embedding function should return `number`.  This will be converted into
-   * an Arrow float array.  By default this will be Float32 but this property can
-   * be used to control the conversion.
+   * @param optionsOrDatatype - The options for the field or the datatype
+   *
+   * @see {@link lancedb.LanceSchema}
    */
-  embeddingDataType?: Float;
+  sourceField(
+    optionsOrDatatype: Partial<FieldOptions> | DataType,
+  ): [DataType, Map<string, EmbeddingFunction>] {
+    let datatype = isDataType(optionsOrDatatype)
+      ? optionsOrDatatype
+      : optionsOrDatatype?.datatype;
+    if (!datatype) {
+      throw new Error("Datatype is required");
+    }
+    datatype = sanitizeType(datatype);
+    const metadata = new Map<string, EmbeddingFunction>();
+    metadata.set("source_column_for", this);
+
+    return [datatype, metadata];
+  }
 
   /**
-   * The dimension of the embedding
+   * vectorField is used in combination with `LanceSchema` to provide a declarative data model
    *
-   * This is optional, normally this can be determined by looking at the results of
-   * `embed`.  If this is not specified, and there is an attempt to apply the embedding
-   * to an empty table, then that process will fail.
+   * @param options - The options for the field
+   *
+   * @see {@link lancedb.LanceSchema}
    */
-  embeddingDimension?: number;
+  vectorField(
+    optionsOrDatatype?: Partial<FieldOptions> | DataType,
+  ): [DataType, Map<string, EmbeddingFunction>] {
+    let dtype: DataType | undefined;
+    let vectorType: DataType;
+    let dims: number | undefined = this.ndims();
 
-  /**
-   * The name of the column that will contain the embedding
-   *
-   * By default this is "vector"
-   */
-  destColumn?: string;
+    // `func.vectorField(new Float32())`
+    if (isDataType(optionsOrDatatype)) {
+      dtype = optionsOrDatatype;
+    } else {
+      // `func.vectorField({
+      //  datatype: new Float32(),
+      //  dims: 10
+      // })`
+      dims = dims ?? optionsOrDatatype?.dims;
+      dtype = optionsOrDatatype?.datatype;
+    }
 
-  /**
-   * Should the source column be excluded from the resulting table
-   *
-   * By default the source column is included.  Set this to true and
-   * only the embedding will be stored.
-   */
-  excludeSource?: boolean;
+    if (dtype !== undefined) {
+      // `func.vectorField(new FixedSizeList(dims, new Field("item", new Float32(), true)))`
+      // or `func.vectorField({datatype: new FixedSizeList(dims, new Field("item", new Float32(), true))})`
+      if (isFixedSizeList(dtype)) {
+        vectorType = dtype;
+        // `func.vectorField(new Float32())`
+        // or `func.vectorField({datatype: new Float32()})`
+      } else if (isFloat(dtype)) {
+        // No `ndims` impl and no `{dims: n}` provided;
+        if (dims === undefined) {
+          throw new Error("ndims is required for vector field");
+        }
+        vectorType = newVectorType(dims, dtype);
+      } else {
+        throw new Error(
+          "Expected FixedSizeList or Float as datatype for vector field",
+        );
+      }
+    } else {
+      if (dims === undefined) {
+        throw new Error("ndims is required for vector field");
+      }
+      vectorType = new FixedSizeList(
+        dims,
+        new Field("item", new Float32(), true),
+      );
+    }
+    const metadata = new Map<string, EmbeddingFunction>();
+    metadata.set("vector_column_for", this);
+
+    return [vectorType, metadata];
+  }
+
+  /** The number of dimensions of the embeddings */
+  ndims(): number | undefined {
+    return undefined;
+  }
+
+  /** The datatype of the embeddings */
+  abstract embeddingDataType(): Float;
 
   /**
    * Creates a vector representation for the given values.
    */
-  embed: (data: T[]) => Promise<number[][]>;
+  abstract computeSourceEmbeddings(
+    data: T[],
+  ): Promise<number[][] | Float32Array[] | Float64Array[]>;
+
+  /**
+  Compute the embeddings for a single query
+ */
+  async computeQueryEmbeddings(data: T): Promise<IntoVector> {
+    return this.computeSourceEmbeddings([data]).then(
+      (embeddings) => embeddings[0],
+    );
+  }
 }
 
-/** Test if the input seems to be an embedding function */
-export function isEmbeddingFunction<T>(
-  value: unknown,
-): value is EmbeddingFunction<T> {
-  if (typeof value !== "object" || value === null) {
-    return false;
-  }
-  if (!("sourceColumn" in value) || !("embed" in value)) {
-    return false;
-  }
-  return (
-    typeof value.sourceColumn === "string" && typeof value.embed === "function"
-  );
+export interface FieldOptions<T extends DataType = DataType> {
+  datatype: T;
+  dims?: number;
 }

nodejs/lancedb/embedding/index.ts

@@ -1,2 +1,113 @@
-export { EmbeddingFunction, isEmbeddingFunction } from "./embedding_function";
-export { OpenAIEmbeddingFunction } from "./openai";
+// Copyright 2023 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { DataType, Field, Schema } from "../arrow";
+import { isDataType } from "../arrow";
+import { sanitizeType } from "../sanitize";
+import { EmbeddingFunction } from "./embedding_function";
+import { EmbeddingFunctionConfig, getRegistry } from "./registry";
+
+export { EmbeddingFunction } from "./embedding_function";
+
+// We need to explicitly export '*' so that the `register` decorator actually registers the class.
+export * from "./openai";
+export * from "./registry";
+
+/**
+ * Create a schema with embedding functions.
+ *
+ * @param fields
+ * @returns Schema
+ * @example
+ * ```ts
+ * class MyEmbeddingFunction extends EmbeddingFunction {
+ * // ...
+ * }
+ * const func = new MyEmbeddingFunction();
+ * const schema = LanceSchema({
+ *   id: new Int32(),
+ *   text: func.sourceField(new Utf8()),
+ *   vector: func.vectorField(),
+ *   // optional: specify the datatype and/or dimensions
+ *   vector2: func.vectorField({ datatype: new Float32(), dims: 3}),
+ * });
+ *
+ * const table = await db.createTable("my_table", data, { schema });
+ * ```
+ */
+export function LanceSchema(
+  fields: Record<string, [object, Map<string, EmbeddingFunction>] | object>,
+): Schema {
+  const arrowFields: Field[] = [];
+
+  const embeddingFunctions = new Map<
+    EmbeddingFunction,
+    Partial<EmbeddingFunctionConfig>
+  >();
+  Object.entries(fields).forEach(([key, value]) => {
+    if (isDataType(value)) {
+      arrowFields.push(new Field(key, sanitizeType(value), true));
+    } else {
+      const [dtype, metadata] = value as [
+        object,
+        Map<string, EmbeddingFunction>,
+      ];
+      arrowFields.push(new Field(key, sanitizeType(dtype), true));
+      parseEmbeddingFunctions(embeddingFunctions, key, metadata);
+    }
+  });
+  const registry = getRegistry();
+  const metadata = registry.getTableMetadata(
+    Array.from(embeddingFunctions.values()) as EmbeddingFunctionConfig[],
+  );
+  const schema = new Schema(arrowFields, metadata);
+  return schema;
+}
+
+function parseEmbeddingFunctions(
+  embeddingFunctions: Map<EmbeddingFunction, Partial<EmbeddingFunctionConfig>>,
+  key: string,
+  metadata: Map<string, EmbeddingFunction>,
+): void {
+  if (metadata.has("source_column_for")) {
+    const embedFunction = metadata.get("source_column_for")!;
+    const current = embeddingFunctions.get(embedFunction);
+    if (current !== undefined) {
+      embeddingFunctions.set(embedFunction, {
+        ...current,
+        sourceColumn: key,
+      });
+    } else {
+      embeddingFunctions.set(embedFunction, {
+        sourceColumn: key,
+        function: embedFunction,
+      });
+    }
+  } else if (metadata.has("vector_column_for")) {
+    const embedFunction = metadata.get("vector_column_for")!;
+
+    const current = embeddingFunctions.get(embedFunction);
+    if (current !== undefined) {
+      embeddingFunctions.set(embedFunction, {
+        ...current,
+        vectorColumn: key,
+      });
+    } else {
+      embeddingFunctions.set(embedFunction, {
+        vectorColumn: key,
+        function: embedFunction,
+      });
+    }
+  }
+}

nodejs/lancedb/embedding/openai.ts

@@ -13,17 +13,31 @@
 // limitations under the License.
 
 import type OpenAI from "openai";
-import { type EmbeddingFunction } from "./embedding_function";
+import { Float, Float32 } from "../arrow";
+import { EmbeddingFunction } from "./embedding_function";
+import { register } from "./registry";
 
-export class OpenAIEmbeddingFunction implements EmbeddingFunction<string> {
-  private readonly _openai: OpenAI;
-  private readonly _modelName: string;
+export type OpenAIOptions = {
+  apiKey?: string;
+  model?: string;
+};
+
+@register("openai")
+export class OpenAIEmbeddingFunction extends EmbeddingFunction<
+  string,
+  OpenAIOptions
+> {
+  #openai: OpenAI;
+  #modelName: string;
+
+  constructor(options: OpenAIOptions = { model: "text-embedding-ada-002" }) {
+    super();
+    const openAIKey = options?.apiKey ?? process.env.OPENAI_API_KEY;
+    if (!openAIKey) {
+      throw new Error("OpenAI API key is required");
+    }
+    const modelName = options?.model ?? "text-embedding-ada-002";
 
-  constructor(
-    sourceColumn: string,
-    openAIKey: string,
-    modelName: string = "text-embedding-ada-002",
-  ) {
     /**
      * @type {import("openai").default}
      */
@@ -36,18 +50,40 @@ export class OpenAIEmbeddingFunction implements EmbeddingFunction<string> {
       throw new Error("please install openai@^4.24.1 using npm install openai");
     }
 
-    this.sourceColumn = sourceColumn;
     const configuration = {
       apiKey: openAIKey,
     };
 
-    this._openai = new Openai(configuration);
-    this._modelName = modelName;
+    this.#openai = new Openai(configuration);
+    this.#modelName = modelName;
   }
 
-  async embed(data: string[]): Promise<number[][]> {
-    const response = await this._openai.embeddings.create({
-      model: this._modelName,
+  toJSON() {
+    return {
+      model: this.#modelName,
+    };
+  }
+
+  ndims(): number {
+    switch (this.#modelName) {
+      case "text-embedding-ada-002":
+        return 1536;
+      case "text-embedding-3-large":
+        return 3072;
+      case "text-embedding-3-small":
+        return 1536;
+      default:
+        return null as never;
+    }
+  }
+
+  embeddingDataType(): Float {
+    return new Float32();
+  }
+
+  async computeSourceEmbeddings(data: string[]): Promise<number[][]> {
+    const response = await this.#openai.embeddings.create({
+      model: this.#modelName,
       input: data,
     });
 
@@ -58,5 +94,15 @@ export class OpenAIEmbeddingFunction implements EmbeddingFunction<string> {
     return embeddings;
   }
 
-  sourceColumn: string;
+  async computeQueryEmbeddings(data: string): Promise<number[]> {
+    if (typeof data !== "string") {
+      throw new Error("Data must be a string");
+    }
+    const response = await this.#openai.embeddings.create({
+      model: this.#modelName,
+      input: data,
+    });
+
+    return response.data[0].embedding;
+  }
 }

nodejs/lancedb/embedding/registry.ts

@@ -0,0 +1,176 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import type { EmbeddingFunction } from "./embedding_function";
+import "reflect-metadata";
+
+export interface EmbeddingFunctionOptions {
+  [key: string]: unknown;
+}
+
+export interface EmbeddingFunctionFactory<
+  T extends EmbeddingFunction = EmbeddingFunction,
+> {
+  new (modelOptions?: EmbeddingFunctionOptions): T;
+}
+
+interface EmbeddingFunctionCreate<T extends EmbeddingFunction> {
+  create(options?: EmbeddingFunctionOptions): T;
+}
+
+/**
+ * This is a singleton class used to register embedding functions
+ * and fetch them by name. It also handles serializing and deserializing.
+ * You can implement your own embedding function by subclassing EmbeddingFunction
+ * or TextEmbeddingFunction and registering it with the registry
+ */
+export class EmbeddingFunctionRegistry {
+  #functions: Map<string, EmbeddingFunctionFactory> = new Map();
+
+  /**
+   * Register an embedding function
+   * @param name The name of the function
+   * @param func The function to register
+   * @throws Error if the function is already registered
+   */
+  register<T extends EmbeddingFunctionFactory = EmbeddingFunctionFactory>(
+    this: EmbeddingFunctionRegistry,
+    alias?: string,
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+  ): (ctor: T) => any {
+    const self = this;
+    return function (ctor: T) {
+      if (!alias) {
+        alias = ctor.name;
+      }
+      if (self.#functions.has(alias)) {
+        throw new Error(
+          `Embedding function with alias "${alias}" already exists`,
+        );
+      }
+      self.#functions.set(alias, ctor);
+      Reflect.defineMetadata("lancedb::embedding::name", alias, ctor);
+      return ctor;
+    };
+  }
+
+  /**
+   * Fetch an embedding function by name
+   * @param name The name of the function
+   */
+  get<T extends EmbeddingFunction<unknown> = EmbeddingFunction>(
+    name: string,
+  ): EmbeddingFunctionCreate<T> | undefined {
+    const factory = this.#functions.get(name);
+    if (!factory) {
+      return undefined;
+    }
+    return {
+      create: function (options: EmbeddingFunctionOptions) {
+        return new factory(options) as unknown as T;
+      },
+    };
+  }
+
+  /**
+   * reset the registry to the initial state
+   */
+  reset(this: EmbeddingFunctionRegistry) {
+    this.#functions.clear();
+  }
+
+  /**
+   * @ignore
+   */
+  parseFunctions(
+    this: EmbeddingFunctionRegistry,
+    metadata: Map<string, string>,
+  ): Map<string, EmbeddingFunctionConfig> {
+    if (!metadata.has("embedding_functions")) {
+      return new Map();
+    } else {
+      type FunctionConfig = {
+        name: string;
+        sourceColumn: string;
+        vectorColumn: string;
+        model: EmbeddingFunctionOptions;
+      };
+      const functions = <FunctionConfig[]>(
+        JSON.parse(metadata.get("embedding_functions")!)
+      );
+      return new Map(
+        functions.map((f) => {
+          const fn = this.get(f.name);
+          if (!fn) {
+            throw new Error(`Function "${f.name}" not found in registry`);
+          }
+          return [
+            f.name,
+            {
+              sourceColumn: f.sourceColumn,
+              vectorColumn: f.vectorColumn,
+              function: this.get(f.name)!.create(f.model),
+            },
+          ];
+        }),
+      );
+    }
+  }
+  // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+  functionToMetadata(conf: EmbeddingFunctionConfig): Record<string, any> {
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    const metadata: Record<string, any> = {};
+    const name = Reflect.getMetadata(
+      "lancedb::embedding::name",
+      conf.function.constructor,
+    );
+    metadata["sourceColumn"] = conf.sourceColumn;
+    metadata["vectorColumn"] = conf.vectorColumn ?? "vector";
+    metadata["name"] = name ?? conf.function.constructor.name;
+    metadata["model"] = conf.function.toJSON();
+    return metadata;
+  }
+
+  getTableMetadata(functions: EmbeddingFunctionConfig[]): Map<string, string> {
+    const metadata = new Map<string, string>();
+    const jsonData = functions.map((conf) => this.functionToMetadata(conf));
+    metadata.set("embedding_functions", JSON.stringify(jsonData));
+
+    return metadata;
+  }
+}
+
+const _REGISTRY = new EmbeddingFunctionRegistry();
+
+export function register(name?: string) {
+  return _REGISTRY.register(name);
+}
+
+/**
+ * Utility function to get the global instance of the registry
+ * @returns `EmbeddingFunctionRegistry` The global instance of the registry
+ * @example
+ * ```ts
+ * const registry = getRegistry();
+ * const openai = registry.get("openai").create();
+ */
+export function getRegistry(): EmbeddingFunctionRegistry {
+  return _REGISTRY;
+}
+
+export interface EmbeddingFunctionConfig {
+  sourceColumn: string;
+  vectorColumn?: string;
+  function: EmbeddingFunction;
+}

nodejs/lancedb/query.ts

@@ -12,7 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Table as ArrowTable, RecordBatch, tableFromIPC } from "apache-arrow";
+import {
+  Table as ArrowTable,
+  type IntoVector,
+  RecordBatch,
+  tableFromIPC,
+} from "./arrow";
 import { type IvfPqOptions } from "./indices";
 import {
   RecordBatchIterator as NativeBatchIterator,
@@ -50,6 +55,39 @@ export class RecordBatchIterator implements AsyncIterator<RecordBatch> {
 }
 /* eslint-enable */
 
+class RecordBatchIterable<
+  NativeQueryType extends NativeQuery | NativeVectorQuery,
+> implements AsyncIterable<RecordBatch>
+{
+  private inner: NativeQueryType;
+  private options?: QueryExecutionOptions;
+
+  constructor(inner: NativeQueryType, options?: QueryExecutionOptions) {
+    this.inner = inner;
+    this.options = options;
+  }
+
+  // biome-ignore lint/suspicious/noExplicitAny: skip
+  [Symbol.asyncIterator](): AsyncIterator<RecordBatch<any>, any, undefined> {
+    return new RecordBatchIterator(
+      this.inner.execute(this.options?.maxBatchLength),
+    );
+  }
+}
+
+/**
+ * Options that control the behavior of a particular query execution
+ */
+export interface QueryExecutionOptions {
+  /**
+   * The maximum number of rows to return in a single batch
+   *
+   * Batches may have fewer rows if the underlying data is stored
+   * in smaller chunks.
+   */
+  maxBatchLength?: number;
+}
+
 /** Common methods supported by all query types */
 export class QueryBase<
   NativeQueryType extends NativeQuery | NativeVectorQuery,
@@ -108,9 +146,12 @@ export class QueryBase<
    * object insertion order is easy to get wrong and `Map` is more foolproof.
    */
   select(
-    columns: string[] | Map<string, string> | Record<string, string>,
+    columns: string[] | Map<string, string> | Record<string, string> | string,
   ): QueryType {
     let columnTuples: [string, string][];
+    if (typeof columns === "string") {
+      columns = [columns];
+    }
     if (Array.isArray(columns)) {
       columnTuples = columns.map((c) => [c, c]);
     } else if (columns instanceof Map) {
@@ -133,8 +174,10 @@ export class QueryBase<
     return this as unknown as QueryType;
   }
 
-  protected nativeExecute(): Promise<NativeBatchIterator> {
-    return this.inner.execute();
+  protected nativeExecute(
+    options?: Partial<QueryExecutionOptions>,
+  ): Promise<NativeBatchIterator> {
+    return this.inner.execute(options?.maxBatchLength);
   }
 
   /**
@@ -148,8 +191,10 @@ export class QueryBase<
    * single query)
    *
    */
-  protected execute(): RecordBatchIterator {
-    return new RecordBatchIterator(this.nativeExecute());
+  protected execute(
+    options?: Partial<QueryExecutionOptions>,
+  ): RecordBatchIterator {
+    return new RecordBatchIterator(this.nativeExecute(options));
   }
 
   // biome-ignore lint/suspicious/noExplicitAny: skip
@@ -159,18 +204,18 @@ export class QueryBase<
   }
 
   /** Collect the results as an Arrow @see {@link ArrowTable}. */
-  async toArrow(): Promise<ArrowTable> {
+  async toArrow(options?: Partial<QueryExecutionOptions>): Promise<ArrowTable> {
     const batches = [];
-    for await (const batch of this) {
+    for await (const batch of new RecordBatchIterable(this.inner, options)) {
       batches.push(batch);
     }
     return new ArrowTable(batches);
   }
 
   /** Collect the results as an array of objects. */
-  async toArray(): Promise<unknown[]> {
-    const tbl = await this.toArrow();
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+  // biome-ignore lint/suspicious/noExplicitAny: arrow.toArrow() returns any[]
+  async toArray(options?: Partial<QueryExecutionOptions>): Promise<any[]> {
+    const tbl = await this.toArrow(options);
     return tbl.toArray();
   }
 }
@@ -369,9 +414,8 @@ export class Query extends QueryBase<NativeQuery, Query> {
    * Vector searches always have a `limit`.  If `limit` has not been called then
    * a default `limit` of 10 will be used.  @see {@link Query#limit}
    */
-  nearestTo(vector: unknown): VectorQuery {
-    // biome-ignore lint/suspicious/noExplicitAny: skip
-    const vectorQuery = this.inner.nearestTo(Float32Array.from(vector as any));
+  nearestTo(vector: IntoVector): VectorQuery {
+    const vectorQuery = this.inner.nearestTo(Float32Array.from(vector));
     return new VectorQuery(vectorQuery);
   }
 }

nodejs/lancedb/sanitize.ts

@@ -20,6 +20,7 @@
 // comes from the exact same library instance.  This is not always the case
 // and so we must sanitize the input to ensure that it is compatible.
 
+import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
 import {
   Binary,
   Bool,
@@ -75,10 +76,9 @@ import {
   Uint64,
   Union,
   Utf8,
-} from "apache-arrow";
-import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
+} from "./arrow";
 
-function sanitizeMetadata(
+export function sanitizeMetadata(
   metadataLike?: unknown,
 ): Map<string, string> | undefined {
   if (metadataLike === undefined || metadataLike === null) {
@@ -97,7 +97,7 @@ function sanitizeMetadata(
   return metadataLike as Map<string, string>;
 }
 
-function sanitizeInt(typeLike: object) {
+export function sanitizeInt(typeLike: object) {
   if (
     !("bitWidth" in typeLike) ||
     typeof typeLike.bitWidth !== "number" ||
@@ -111,14 +111,14 @@ function sanitizeInt(typeLike: object) {
   return new Int(typeLike.isSigned, typeLike.bitWidth as IntBitWidth);
 }
 
-function sanitizeFloat(typeLike: object) {
+export function sanitizeFloat(typeLike: object) {
   if (!("precision" in typeLike) || typeof typeLike.precision !== "number") {
     throw Error("Expected a Float Type to have a `precision` property");
   }
   return new Float(typeLike.precision as Precision);
 }
 
-function sanitizeDecimal(typeLike: object) {
+export function sanitizeDecimal(typeLike: object) {
   if (
     !("scale" in typeLike) ||
     typeof typeLike.scale !== "number" ||
@@ -134,14 +134,14 @@ function sanitizeDecimal(typeLike: object) {
   return new Decimal(typeLike.scale, typeLike.precision, typeLike.bitWidth);
 }
 
-function sanitizeDate(typeLike: object) {
+export function sanitizeDate(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Date type to have a `unit` property");
   }
   return new Date_(typeLike.unit as DateUnit);
 }
 
-function sanitizeTime(typeLike: object) {
+export function sanitizeTime(typeLike: object) {
   if (
     !("unit" in typeLike) ||
     typeof typeLike.unit !== "number" ||
@@ -155,7 +155,7 @@ function sanitizeTime(typeLike: object) {
   return new Time(typeLike.unit, typeLike.bitWidth as TimeBitWidth);
 }
 
-function sanitizeTimestamp(typeLike: object) {
+export function sanitizeTimestamp(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Timestamp type to have a `unit` property");
   }
@@ -166,7 +166,7 @@ function sanitizeTimestamp(typeLike: object) {
   return new Timestamp(typeLike.unit, timezone);
 }
 
-function sanitizeTypedTimestamp(
+export function sanitizeTypedTimestamp(
   typeLike: object,
   // eslint-disable-next-line @typescript-eslint/naming-convention
   Datatype:
@@ -182,14 +182,14 @@ function sanitizeTypedTimestamp(
   return new Datatype(timezone);
 }
 
-function sanitizeInterval(typeLike: object) {
+export function sanitizeInterval(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected an Interval type to have a `unit` property");
   }
   return new Interval(typeLike.unit);
 }
 
-function sanitizeList(typeLike: object) {
+export function sanitizeList(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a List type to have an array-like `children` property",
@@ -201,7 +201,7 @@ function sanitizeList(typeLike: object) {
   return new List(sanitizeField(typeLike.children[0]));
 }
 
-function sanitizeStruct(typeLike: object) {
+export function sanitizeStruct(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a Struct type to have an array-like `children` property",
@@ -210,7 +210,7 @@ function sanitizeStruct(typeLike: object) {
   return new Struct(typeLike.children.map((child) => sanitizeField(child)));
 }
 
-function sanitizeUnion(typeLike: object) {
+export function sanitizeUnion(typeLike: object) {
   if (
     !("typeIds" in typeLike) ||
     !("mode" in typeLike) ||
@@ -234,7 +234,7 @@ function sanitizeUnion(typeLike: object) {
   );
 }
 
-function sanitizeTypedUnion(
+export function sanitizeTypedUnion(
   typeLike: object,
   // eslint-disable-next-line @typescript-eslint/naming-convention
   UnionType: typeof DenseUnion | typeof SparseUnion,
@@ -256,7 +256,7 @@ function sanitizeTypedUnion(
   );
 }
 
-function sanitizeFixedSizeBinary(typeLike: object) {
+export function sanitizeFixedSizeBinary(typeLike: object) {
   if (!("byteWidth" in typeLike) || typeof typeLike.byteWidth !== "number") {
     throw Error(
       "Expected a FixedSizeBinary type to have a `byteWidth` property",
@@ -265,7 +265,7 @@ function sanitizeFixedSizeBinary(typeLike: object) {
   return new FixedSizeBinary(typeLike.byteWidth);
 }
 
-function sanitizeFixedSizeList(typeLike: object) {
+export function sanitizeFixedSizeList(typeLike: object) {
   if (!("listSize" in typeLike) || typeof typeLike.listSize !== "number") {
     throw Error("Expected a FixedSizeList type to have a `listSize` property");
   }
@@ -283,7 +283,7 @@ function sanitizeFixedSizeList(typeLike: object) {
   );
 }
 
-function sanitizeMap(typeLike: object) {
+export function sanitizeMap(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a Map type to have an array-like `children` property",
@@ -300,14 +300,14 @@ function sanitizeMap(typeLike: object) {
   );
 }
 
-function sanitizeDuration(typeLike: object) {
+export function sanitizeDuration(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Duration type to have a `unit` property");
   }
   return new Duration(typeLike.unit);
 }
 
-function sanitizeDictionary(typeLike: object) {
+export function sanitizeDictionary(typeLike: object) {
   if (!("id" in typeLike) || typeof typeLike.id !== "number") {
     throw Error("Expected a Dictionary type to have an `id` property");
   }
@@ -329,7 +329,7 @@ function sanitizeDictionary(typeLike: object) {
 }
 
 // biome-ignore lint/suspicious/noExplicitAny: skip
-function sanitizeType(typeLike: unknown): DataType<any> {
+export function sanitizeType(typeLike: unknown): DataType<any> {
   if (typeof typeLike !== "object" || typeLike === null) {
     throw Error("Expected a Type but object was null/undefined");
   }
@@ -449,7 +449,7 @@ function sanitizeType(typeLike: unknown): DataType<any> {
   }
 }
 
-function sanitizeField(fieldLike: unknown): Field {
+export function sanitizeField(fieldLike: unknown): Field {
   if (fieldLike instanceof Field) {
     return fieldLike;
   }

nodejs/lancedb/table.ts

@@ -12,18 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Schema, tableFromIPC } from "apache-arrow";
-import { Data, fromDataToBuffer } from "./arrow";
+import {
+  Table as ArrowTable,
+  Data,
+  IntoVector,
+  Schema,
+  fromDataToBuffer,
+  tableFromIPC,
+} from "./arrow";
+
+import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
 import { IndexOptions } from "./indices";
 import {
   AddColumnsSql,
   ColumnAlteration,
   IndexConfig,
+  OptimizeStats,
   Table as _NativeTable,
 } from "./native";
 import { Query, VectorQuery } from "./query";
-
 export { IndexConfig } from "./native";
+
 /**
  * Options for adding data to a table.
  */
@@ -50,6 +59,23 @@ export interface UpdateOptions {
   where: string;
 }
 
+export interface OptimizeOptions {
+  /**
+   * If set then all versions older than the given date
+   * be removed.  The current version will never be removed.
+   * The default is 7 days
+   * @example
+   * // Delete all versions older than 1 day
+   * const olderThan = new Date();
+   * olderThan.setDate(olderThan.getDate() - 1));
+   * tbl.cleanupOlderVersions(olderThan);
+   *
+   * // Delete all versions except the current version
+   * tbl.cleanupOlderVersions(new Date());
+   */
+  cleanupOlderThan: Date;
+}
+
 /**
  * A Table is a collection of Records in a LanceDB Database.
  *
@@ -91,6 +117,14 @@ export class Table {
     return this.inner.display();
   }
 
+  async #getEmbeddingFunctions(): Promise<
+    Map<string, EmbeddingFunctionConfig>
+  > {
+    const schema = await this.schema();
+    const registry = getRegistry();
+    return registry.parseFunctions(schema.metadata);
+  }
+
   /** Get the schema of the table. */
   async schema(): Promise<Schema> {
     const schemaBuf = await this.inner.schema();
@@ -104,8 +138,15 @@ export class Table {
    */
   async add(data: Data, options?: Partial<AddDataOptions>): Promise<void> {
     const mode = options?.mode ?? "append";
+    const schema = await this.schema();
+    const registry = getRegistry();
+    const functions = registry.parseFunctions(schema.metadata);
 
-    const buffer = await fromDataToBuffer(data);
+    const buffer = await fromDataToBuffer(
+      data,
+      functions.values().next().value,
+      schema,
+    );
     await this.inner.add(buffer, mode);
   }
 
@@ -245,6 +286,40 @@ export class Table {
     return new Query(this.inner);
   }
 
+  /**
+   * Create a search query to find the nearest neighbors
+   * of the given query vector
+   * @param {string} query - the query. This will be converted to a vector using the table's provided embedding function
+   * @rejects {Error} If no embedding functions are defined in the table
+   */
+  search(query: string): Promise<VectorQuery>;
+  /**
+   * Create a search query to find the nearest neighbors
+   * of the given query vector
+   * @param {IntoVector} query - the query vector
+   */
+  search(query: IntoVector): VectorQuery;
+  search(query: string | IntoVector): Promise<VectorQuery> | VectorQuery {
+    if (typeof query !== "string") {
+      return this.vectorSearch(query);
+    } else {
+      return this.#getEmbeddingFunctions().then(async (functions) => {
+        // TODO: Support multiple embedding functions
+        const embeddingFunc: EmbeddingFunctionConfig | undefined = functions
+          .values()
+          .next().value;
+        if (!embeddingFunc) {
+          return Promise.reject(
+            new Error("No embedding functions are defined in the table"),
+          );
+        }
+        const embeddings =
+          await embeddingFunc.function.computeQueryEmbeddings(query);
+        return this.query().nearestTo(embeddings);
+      });
+    }
+  }
+
   /**
    * Search the table with a given query vector.
    *
@@ -252,7 +327,7 @@ export class Table {
    * is the same thing as calling `nearestTo` on the builder returned
    * by `query`.  @see {@link Query#nearestTo} for more details.
    */
-  vectorSearch(vector: unknown): VectorQuery {
+  vectorSearch(vector: IntoVector): VectorQuery {
     return this.query().nearestTo(vector);
   }
 
@@ -352,8 +427,55 @@ export class Table {
     await this.inner.restore();
   }
 
+  /**
+   * Optimize the on-disk data and indices for better performance.
+   *
+   * Modeled after ``VACUUM`` in PostgreSQL.
+   *
+   *  Optimization covers three operations:
+   *
+   *  - Compaction: Merges small files into larger ones
+   *  - Prune: Removes old versions of the dataset
+   *  - Index: Optimizes the indices, adding new data to existing indices
+   *
+   *
+   *  Experimental API
+   *  ----------------
+   *
+   *  The optimization process is undergoing active development and may change.
+   *  Our goal with these changes is to improve the performance of optimization and
+   *  reduce the complexity.
+   *
+   *  That being said, it is essential today to run optimize if you want the best
+   *  performance.  It should be stable and safe to use in production, but it our
+   *  hope that the API may be simplified (or not even need to be called) in the
+   *  future.
+   *
+   *  The frequency an application shoudl call optimize is based on the frequency of
+   *  data modifications.  If data is frequently added, deleted, or updated then
+   *  optimize should be run frequently.  A good rule of thumb is to run optimize if
+   *  you have added or modified 100,000 or more records or run more than 20 data
+   *  modification operations.
+   */
+  async optimize(options?: Partial<OptimizeOptions>): Promise<OptimizeStats> {
+    let cleanupOlderThanMs;
+    if (
+      options?.cleanupOlderThan !== undefined &&
+      options?.cleanupOlderThan !== null
+    ) {
+      cleanupOlderThanMs =
+        new Date().getTime() - options.cleanupOlderThan.getTime();
+    }
+    return await this.inner.optimize(cleanupOlderThanMs);
+  }
+
   /** List all indices that have been created with {@link Table.createIndex} */
   async listIndices(): Promise<IndexConfig[]> {
     return await this.inner.listIndices();
   }
+
+  /** Return the table as an arrow table */
+  async toArrow(): Promise<ArrowTable> {
+    return await this.query().toArrow();
+  }
 }

nodejs/npm/darwin-arm64/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@lancedb/lancedb-darwin-arm64",
-	"version": "0.4.20",
+	"version": "0.5.1",
 	"os": ["darwin"],
 	"cpu": ["arm64"],
 	"main": "lancedb.darwin-arm64.node",

nodejs/npm/darwin-x64/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@lancedb/lancedb-darwin-x64",
-	"version": "0.4.20",
+	"version": "0.5.1",
 	"os": ["darwin"],
 	"cpu": ["x64"],
 	"main": "lancedb.darwin-x64.node",

nodejs/npm/linux-arm64-gnu/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@lancedb/lancedb-linux-arm64-gnu",
-	"version": "0.4.20",
+	"version": "0.5.1",
 	"os": ["linux"],
 	"cpu": ["arm64"],
 	"main": "lancedb.linux-arm64-gnu.node",

nodejs/npm/linux-x64-gnu/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@lancedb/lancedb-linux-x64-gnu",
-	"version": "0.4.20",
+	"version": "0.5.1",
 	"os": ["linux"],
 	"cpu": ["x64"],
 	"main": "lancedb.linux-x64-gnu.node",

nodejs/npm/win32-x64-msvc/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@lancedb/lancedb-win32-x64-msvc",
-	"version": "0.4.14",
+	"version": "0.5.1",
 	"os": ["win32"],
 	"cpu": ["x64"],
 	"main": "lancedb.win32-x64-msvc.node",

nodejs/package.json

@@ -1,8 +1,12 @@
 {
   "name": "@lancedb/lancedb",
-  "version": "0.4.20",
-  "main": "./dist/index.js",
-  "types": "./dist/index.d.ts",
+  "version": "0.5.1",
+  "main": "dist/index.js",
+  "exports": {
+    ".": "./dist/index.js",
+    "./embedding": "./dist/embedding/index.js"
+  },
+  "types": "dist/index.d.ts",
   "napi": {
     "name": "lancedb",
     "triples": {
@@ -62,6 +66,7 @@
   },
   "dependencies": {
     "apache-arrow": "^15.0.0",
-    "openai": "^4.29.2"
+    "openai": "^4.29.2",
+    "reflect-metadata": "^0.2.2"
   }
 }

nodejs/src/connection.rs

@@ -126,6 +126,7 @@ impl Connection {
         buf: Buffer,
         mode: String,
         storage_options: Option<HashMap<String, String>>,
+        use_legacy_format: Option<bool>,
     ) -> napi::Result<Table> {
         let batches = ipc_file_to_batches(buf.to_vec())
             .map_err(|e| napi::Error::from_reason(format!("Failed to read IPC file: {}", e)))?;
@@ -136,6 +137,9 @@ impl Connection {
                 builder = builder.storage_option(key, value);
             }
         }
+        if let Some(use_legacy_format) = use_legacy_format {
+            builder = builder.use_legacy_format(use_legacy_format);
+        }
         let tbl = builder
             .execute()
             .await
@@ -150,6 +154,7 @@ impl Connection {
         schema_buf: Buffer,
         mode: String,
         storage_options: Option<HashMap<String, String>>,
+        use_legacy_format: Option<bool>,
     ) -> napi::Result<Table> {
         let schema = ipc_file_to_schema(schema_buf.to_vec()).map_err(|e| {
             napi::Error::from_reason(format!("Failed to marshal schema from JS to Rust: {}", e))
@@ -164,6 +169,9 @@ impl Connection {
                 builder = builder.storage_option(key, value);
             }
         }
+        if let Some(use_legacy_format) = use_legacy_format {
+            builder = builder.use_legacy_format(use_legacy_format);
+        }
         let tbl = builder
             .execute()
             .await

nodejs/src/lib.rs

@@ -56,6 +56,7 @@ pub enum WriteMode {
 /// Write options when creating a Table.
 #[napi(object)]
 pub struct WriteOptions {
+    /// Write mode for writing to a table.
     pub mode: Option<WriteMode>,
 }
 

nodejs/src/query.rs

@@ -15,6 +15,7 @@
 use lancedb::query::ExecutableQuery;
 use lancedb::query::Query as LanceDbQuery;
 use lancedb::query::QueryBase;
+use lancedb::query::QueryExecutionOptions;
 use lancedb::query::Select;
 use lancedb::query::VectorQuery as LanceDbVectorQuery;
 use napi::bindgen_prelude::*;
@@ -62,10 +63,21 @@ impl Query {
     }
 
     #[napi]
-    pub async fn execute(&self) -> napi::Result<RecordBatchIterator> {
-        let inner_stream = self.inner.execute().await.map_err(|e| {
-            napi::Error::from_reason(format!("Failed to execute query stream: {}", e))
-        })?;
+    pub async fn execute(
+        &self,
+        max_batch_length: Option<u32>,
+    ) -> napi::Result<RecordBatchIterator> {
+        let mut execution_opts = QueryExecutionOptions::default();
+        if let Some(max_batch_length) = max_batch_length {
+            execution_opts.max_batch_length = max_batch_length;
+        }
+        let inner_stream = self
+            .inner
+            .execute_with_options(execution_opts)
+            .await
+            .map_err(|e| {
+                napi::Error::from_reason(format!("Failed to execute query stream: {}", e))
+            })?;
         Ok(RecordBatchIterator::new(inner_stream))
     }
 }
@@ -125,10 +137,21 @@ impl VectorQuery {
     }
 
     #[napi]
-    pub async fn execute(&self) -> napi::Result<RecordBatchIterator> {
-        let inner_stream = self.inner.execute().await.map_err(|e| {
-            napi::Error::from_reason(format!("Failed to execute query stream: {}", e))
-        })?;
+    pub async fn execute(
+        &self,
+        max_batch_length: Option<u32>,
+    ) -> napi::Result<RecordBatchIterator> {
+        let mut execution_opts = QueryExecutionOptions::default();
+        if let Some(max_batch_length) = max_batch_length {
+            execution_opts.max_batch_length = max_batch_length;
+        }
+        let inner_stream = self
+            .inner
+            .execute_with_options(execution_opts)
+            .await
+            .map_err(|e| {
+                napi::Error::from_reason(format!("Failed to execute query stream: {}", e))
+            })?;
         Ok(RecordBatchIterator::new(inner_stream))
     }
 }

nodejs/src/table.rs

@@ -15,8 +15,8 @@
 use arrow_ipc::writer::FileWriter;
 use lancedb::ipc::ipc_file_to_batches;
 use lancedb::table::{
-    AddDataMode, ColumnAlteration as LanceColumnAlteration, NewColumnTransform,
-    Table as LanceDbTable,
+    AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration, NewColumnTransform,
+    OptimizeAction, OptimizeOptions, Table as LanceDbTable,
 };
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
@@ -263,6 +263,60 @@ impl Table {
         self.inner_ref()?.restore().await.default_error()
     }
 
+    #[napi]
+    pub async fn optimize(&self, older_than_ms: Option<i64>) -> napi::Result<OptimizeStats> {
+        let inner = self.inner_ref()?;
+
+        let older_than = if let Some(ms) = older_than_ms {
+            if ms == i64::MIN {
+                return Err(napi::Error::from_reason(format!(
+                    "older_than_ms can not be {}",
+                    i32::MIN,
+                )));
+            }
+            Duration::try_milliseconds(ms)
+        } else {
+            None
+        };
+
+        let compaction_stats = inner
+            .optimize(OptimizeAction::Compact {
+                options: lancedb::table::CompactionOptions::default(),
+                remap_options: None,
+            })
+            .await
+            .default_error()?
+            .compaction
+            .unwrap();
+        let prune_stats = inner
+            .optimize(OptimizeAction::Prune {
+                older_than,
+                delete_unverified: None,
+            })
+            .await
+            .default_error()?
+            .prune
+            .unwrap();
+        inner
+            .optimize(lancedb::table::OptimizeAction::Index(
+                OptimizeOptions::default(),
+            ))
+            .await
+            .default_error()?;
+        Ok(OptimizeStats {
+            compaction: CompactionStats {
+                files_added: compaction_stats.files_added as i64,
+                files_removed: compaction_stats.files_removed as i64,
+                fragments_added: compaction_stats.fragments_added as i64,
+                fragments_removed: compaction_stats.fragments_removed as i64,
+            },
+            prune: RemovalStats {
+                bytes_removed: prune_stats.bytes_removed as i64,
+                old_versions_removed: prune_stats.old_versions as i64,
+            },
+        })
+    }
+
     #[napi]
     pub async fn list_indices(&self) -> napi::Result<Vec<IndexConfig>> {
         Ok(self
@@ -298,6 +352,40 @@ impl From<lancedb::index::IndexConfig> for IndexConfig {
     }
 }
 
+/// Statistics about a compaction operation.
+#[napi(object)]
+#[derive(Clone, Debug)]
+pub struct CompactionStats {
+    /// The number of fragments removed
+    pub fragments_removed: i64,
+    /// The number of new, compacted fragments added
+    pub fragments_added: i64,
+    /// The number of data files removed
+    pub files_removed: i64,
+    /// The number of new, compacted data files added
+    pub files_added: i64,
+}
+
+/// Statistics about a cleanup operation
+#[napi(object)]
+#[derive(Clone, Debug)]
+pub struct RemovalStats {
+    /// The number of bytes removed
+    pub bytes_removed: i64,
+    /// The number of old versions removed
+    pub old_versions_removed: i64,
+}
+
+/// Statistics about an optimize operation
+#[napi(object)]
+#[derive(Clone, Debug)]
+pub struct OptimizeStats {
+    /// Statistics about the compaction operation
+    pub compaction: CompactionStats,
+    /// Statistics about the removal operation
+    pub prune: RemovalStats,
+}
+
 ///  A definition of a column alteration. The alteration changes the column at
 /// `path` to have the new name `name`, to be nullable if `nullable` is true,
 /// and to have the data type `data_type`. At least one of `rename` or `nullable`

nodejs/tsconfig.json

@@ -7,7 +7,9 @@
     "outDir": "./dist",
     "strict": true,
     "allowJs": true,
-    "resolveJsonModule": true
+    "resolveJsonModule": true,
+    "emitDecoratorMetadata": true,
+    "experimentalDecorators": true
   },
   "exclude": ["./dist/*"],
   "typedocOptions": {

python/.bumpversion.cfg

@@ -1,8 +0,0 @@
-[bumpversion]
-current_version = 0.6.13
-commit = True
-message = [python] Bump version: {current_version} → {new_version}
-tag = True
-tag_name = python-v{new_version}
-
-[bumpversion:file:pyproject.toml]

python/.bumpversion.toml

@@ -0,0 +1,34 @@
+[tool.bumpversion]
+current_version = "0.8.2"
+parse = """(?x)
+    (?P<major>0|[1-9]\\d*)\\.
+    (?P<minor>0|[1-9]\\d*)\\.
+    (?P<patch>0|[1-9]\\d*)
+    (?:-(?P<pre_l>[a-zA-Z-]+)\\.(?P<pre_n>0|[1-9]\\d*))?
+"""
+serialize = [
+    "{major}.{minor}.{patch}-{pre_l}.{pre_n}",
+    "{major}.{minor}.{patch}",
+]
+search = "{current_version}"
+replace = "{new_version}"
+regex = false
+ignore_missing_version = false
+ignore_missing_files = false
+tag = true
+sign_tags = false
+tag_name = "python-v{new_version}"
+tag_message = "Bump version: {current_version} → {new_version}"
+allow_dirty = true
+commit = true
+message = "Bump version: {current_version} → {new_version}"
+commit_args = ""
+
+[tool.bumpversion.parts.pre_l]
+values = ["beta", "final"]
+optional_value = "final"
+
+[[tool.bumpversion.files]]
+filename = "Cargo.toml"
+search = "\nversion = \"{current_version}\""
+replace = "\nversion = \"{new_version}\""

python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lancedb-python"
-version = "0.4.10"
+version = "0.8.2"
 edition.workspace = true
 description = "Python bindings for LanceDB"
 license.workspace = true

python/pyproject.toml

@@ -1,16 +1,16 @@
 [project]
 name = "lancedb"
-version = "0.6.13"
+# version in Cargo.toml
 dependencies = [
     "deprecation",
-    "pylance==0.10.12",
+    "pylance==0.12.1",
     "ratelimiter~=1.0",
     "requests>=2.31.0",
     "retry>=0.9.2",
     "tqdm>=4.27.0",
     "pydantic>=1.10",
     "attrs>=21.3.0",
-    "semver",
+    "packaging",
     "cachetools",
     "overrides>=0.7",
 ]

python/python/lancedb/_lancedb.pyi

@@ -24,6 +24,7 @@ class Connection(object):
         mode: str,
         data: pa.RecordBatchReader,
         storage_options: Optional[Dict[str, str]] = None,
+        use_legacy_format: Optional[bool] = None,
     ) -> Table: ...
     async def create_empty_table(
         self,
@@ -31,6 +32,7 @@ class Connection(object):
         mode: str,
         schema: pa.Schema,
         storage_options: Optional[Dict[str, str]] = None,
+        use_legacy_format: Optional[bool] = None,
     ) -> Table: ...
 
 class Table:
@@ -72,7 +74,7 @@ class Query:
     def select(self, columns: Tuple[str, str]): ...
     def limit(self, limit: int): ...
     def nearest_to(self, query_vec: pa.Array) -> VectorQuery: ...
-    async def execute(self) -> RecordBatchStream: ...
+    async def execute(self, max_batch_legnth: Optional[int]) -> RecordBatchStream: ...
 
 class VectorQuery:
     async def execute(self) -> RecordBatchStream: ...
@@ -86,3 +88,17 @@ class VectorQuery:
     def refine_factor(self, refine_factor: int): ...
     def nprobes(self, nprobes: int): ...
     def bypass_vector_index(self): ...
+
+class CompactionStats:
+    fragments_removed: int
+    fragments_added: int
+    files_removed: int
+    files_added: int
+
+class RemovalStats:
+    bytes_removed: int
+    old_versions_removed: int
+
+class OptimizeStats:
+    compaction: CompactionStats
+    prune: RemovalStats

python/python/lancedb/db.py

@@ -509,7 +509,7 @@ class AsyncConnection(object):
         return self._inner.__repr__()
 
     def __enter__(self):
-        self
+        return self
 
     def __exit__(self, *_):
         self.close()
@@ -558,6 +558,8 @@ class AsyncConnection(object):
         on_bad_vectors: Optional[str] = None,
         fill_value: Optional[float] = None,
         storage_options: Optional[Dict[str, str]] = None,
+        *,
+        use_legacy_format: Optional[bool] = None,
     ) -> AsyncTable:
         """Create an [AsyncTable][lancedb.table.AsyncTable] in the database.
 
@@ -600,6 +602,9 @@ class AsyncConnection(object):
             connection will be inherited by the table, but can be overridden here.
             See available options at
             https://lancedb.github.io/lancedb/guides/storage/
+        use_legacy_format: bool, optional, default True
+            If True, use the legacy format for the table. If False, use the new format.
+            The default is True while the new format is in beta.
 
 
         Returns
@@ -761,7 +766,11 @@ class AsyncConnection(object):
 
         if data is None:
             new_table = await self._inner.create_empty_table(
-                name, mode, schema, storage_options=storage_options
+                name,
+                mode,
+                schema,
+                storage_options=storage_options,
+                use_legacy_format=use_legacy_format,
             )
         else:
             data = data_to_reader(data, schema)
@@ -770,6 +779,7 @@ class AsyncConnection(object):
                 mode,
                 data,
                 storage_options=storage_options,
+                use_legacy_format=use_legacy_format,
             )
 
         return AsyncTable(new_table)
@@ -779,7 +789,7 @@ class AsyncConnection(object):
         name: str,
         storage_options: Optional[Dict[str, str]] = None,
         index_cache_size: Optional[int] = None,
-    ) -> Table:
+    ) -> AsyncTable:
         """Open a Lance Table in the database.
 
         Parameters

python/python/lancedb/embeddings/base.py

@@ -153,7 +153,7 @@ class TextEmbeddingFunction(EmbeddingFunction):
 
     @abstractmethod
     def generate_embeddings(
-        self, texts: Union[List[str], np.ndarray]
+        self, texts: Union[List[str], np.ndarray], *args, **kwargs
     ) -> List[np.array]:
         """
         Generate the embeddings for the given texts

python/python/lancedb/embeddings/bedrock.py

@@ -73,8 +73,10 @@ class BedRockText(TextEmbeddingFunction):
     assumed_role: Union[str, None] = None
     profile_name: Union[str, None] = None
     role_session_name: str = "lancedb-embeddings"
+    source_input_type: str = "search_document"
+    query_input_type: str = "search_query"
 
-    if PYDANTIC_VERSION < (2, 0):  # Pydantic 1.x compat
+    if PYDANTIC_VERSION.major < 2:  # Pydantic 1.x compat
 
         class Config:
             keep_untouched = (cached_property,)
@@ -87,21 +89,29 @@ class BedRockText(TextEmbeddingFunction):
         # TODO: fix hardcoding
         if self.name == "amazon.titan-embed-text-v1":
             return 1536
-        elif self.name in {"cohere.embed-english-v3", "cohere.embed-multilingual-v3"}:
+        elif self.name in [
+            "amazon.titan-embed-text-v2:0",
+            "cohere.embed-english-v3",
+            "cohere.embed-multilingual-v3",
+        ]:
+            # TODO: "amazon.titan-embed-text-v2:0" model supports dynamic ndims
             return 1024
         else:
-            raise ValueError(f"Unknown model name: {self.name}")
+            raise ValueError(f"Model {self.name} not supported")
 
     def compute_query_embeddings(
         self, query: str, *args, **kwargs
     ) -> List[List[float]]:
-        return self.compute_source_embeddings(query)
+        return self.compute_source_embeddings(query, input_type=self.query_input_type)
 
     def compute_source_embeddings(
         self, texts: TEXT, *args, **kwargs
     ) -> List[List[float]]:
         texts = self.sanitize_input(texts)
-        return self.generate_embeddings(texts)
+        # assume source input type if not passed by `compute_query_embeddings`
+        kwargs["input_type"] = kwargs.get("input_type") or self.source_input_type
+
+        return self.generate_embeddings(texts, **kwargs)
 
     def generate_embeddings(
         self, texts: Union[List[str], np.ndarray], *args, **kwargs
@@ -121,11 +131,11 @@ class BedRockText(TextEmbeddingFunction):
         """
         results = []
         for text in texts:
-            response = self._generate_embedding(text)
+            response = self._generate_embedding(text, *args, **kwargs)
             results.append(response)
         return results
 
-    def _generate_embedding(self, text: str) -> List[float]:
+    def _generate_embedding(self, text: str, *args, **kwargs) -> List[float]:
         """
         Get the embeddings for the given texts
 
@@ -141,14 +151,12 @@ class BedRockText(TextEmbeddingFunction):
         """
         # format input body for provider
         provider = self.name.split(".")[0]
-        _model_kwargs = {}
-        input_body = {**_model_kwargs}
+        input_body = {**kwargs}
         if provider == "cohere":
-            if "input_type" not in input_body.keys():
-                input_body["input_type"] = "search_document"
             input_body["texts"] = [text]
         else:
             # includes common provider == "amazon"
+            input_body.pop("input_type", None)
             input_body["inputText"] = text
         body = json.dumps(input_body)
 

python/python/lancedb/embeddings/cohere.py

@@ -19,7 +19,7 @@ import numpy as np
 from ..util import attempt_import_or_raise
 from .base import TextEmbeddingFunction
 from .registry import register
-from .utils import api_key_not_found_help
+from .utils import api_key_not_found_help, TEXT
 
 
 @register("cohere")
@@ -32,8 +32,36 @@ class CohereEmbeddingFunction(TextEmbeddingFunction):
     Parameters
     ----------
     name: str, default "embed-multilingual-v2.0"
-        The name of the model to use. See the Cohere documentation for
-        a list of available models.
+        The name of the model to use. List of acceptable models:
+
+            * embed-english-v3.0
+            * embed-multilingual-v3.0
+            * embed-english-light-v3.0
+            * embed-multilingual-light-v3.0
+            * embed-english-v2.0
+            * embed-english-light-v2.0
+            * embed-multilingual-v2.0
+
+    source_input_type: str, default "search_document"
+        The input type for the source column in the database
+
+    query_input_type: str, default "search_query"
+        The input type for the query column in the database
+
+    Cohere supports following input types:
+
+    | Input Type               | Description                          |
+    |-------------------------|---------------------------------------|
+    | "`search_document`"     | Used for embeddings stored in a vector|
+    |                         | database for search use-cases.        |
+    | "`search_query`"        | Used for embeddings of search queries |
+    |                         | run against a vector DB               |
+    | "`semantic_similarity`" | Specifies the given text will be used |
+    |                         | for Semantic Textual Similarity (STS) |
+    | "`classification`"      | Used for embeddings passed through a  |
+    |                         | text classifier.                      |
+    | "`clustering`"          | Used for the embeddings run through a |
+    |                         | clustering algorithm                  |
 
     Examples
     --------
@@ -61,14 +89,39 @@ class CohereEmbeddingFunction(TextEmbeddingFunction):
     """
 
     name: str = "embed-multilingual-v2.0"
+    source_input_type: str = "search_document"
+    query_input_type: str = "search_query"
     client: ClassVar = None
 
     def ndims(self):
         # TODO: fix hardcoding
-        return 768
+        if self.name in [
+            "embed-english-v3.0",
+            "embed-multilingual-v3.0",
+            "embed-english-light-v2.0",
+        ]:
+            return 1024
+        elif self.name in ["embed-english-light-v3.0", "embed-multilingual-light-v3.0"]:
+            return 384
+        elif self.name == "embed-english-v2.0":
+            return 4096
+        elif self.name == "embed-multilingual-v2.0":
+            return 768
+        else:
+            raise ValueError(f"Model {self.name} not supported")
+
+    def compute_query_embeddings(self, query: str, *args, **kwargs) -> List[np.array]:
+        return self.compute_source_embeddings(query, input_type=self.query_input_type)
+
+    def compute_source_embeddings(self, texts: TEXT, *args, **kwargs) -> List[np.array]:
+        texts = self.sanitize_input(texts)
+        input_type = (
+            kwargs.get("input_type") or self.source_input_type
+        )  # assume source input type if not passed by `compute_query_embeddings`
+        return self.generate_embeddings(texts, input_type=input_type)
 
     def generate_embeddings(
-        self, texts: Union[List[str], np.ndarray]
+        self, texts: Union[List[str], np.ndarray], *args, **kwargs
     ) -> List[np.array]:
         """
         Get the embeddings for the given texts
@@ -78,9 +131,10 @@ class CohereEmbeddingFunction(TextEmbeddingFunction):
         texts: list[str] or np.ndarray (of str)
             The texts to embed
         """
-        # TODO retry, rate limit, token limit
         self._init_client()
-        rs = CohereEmbeddingFunction.client.embed(texts=texts, model=self.name)
+        rs = CohereEmbeddingFunction.client.embed(
+            texts=texts, model=self.name, **kwargs
+        )
 
         return [emb for emb in rs.embeddings]
 

python/python/lancedb/embeddings/gemini_text.py

@@ -90,7 +90,7 @@ class GeminiText(TextEmbeddingFunction):
     query_task_type: str = "retrieval_query"
     source_task_type: str = "retrieval_document"
 
-    if PYDANTIC_VERSION < (2, 0):  # Pydantic 1.x compat
+    if PYDANTIC_VERSION.major < 2:  # Pydantic 1.x compat
 
         class Config:
             keep_untouched = (cached_property,)

python/python/lancedb/embeddings/imagebind.py

@@ -40,7 +40,7 @@ class ImageBindEmbeddings(EmbeddingFunction):
     device: str = "cpu"
     normalize: bool = False
 
-    if PYDANTIC_VERSION < (2, 0):  # Pydantic 1.x compat
+    if PYDANTIC_VERSION.major < 2:  # Pydantic 1.x compat
 
         class Config:
             keep_untouched = (cached_property,)

python/python/lancedb/embeddings/transformers.py

@@ -54,7 +54,7 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
         self._tokenizer = transformers.AutoTokenizer.from_pretrained(self.name)
         self._model = transformers.AutoModel.from_pretrained(self.name)
 
-    if PYDANTIC_VERSION < (2, 0):  # Pydantic 1.x compat
+    if PYDANTIC_VERSION.major < 2:  # Pydantic 1.x compat
 
         class Config:
             keep_untouched = (cached_property,)

python/python/lancedb/pydantic.py

@@ -35,13 +35,13 @@ from typing import (
 import numpy as np
 import pyarrow as pa
 import pydantic
-import semver
+from packaging.version import Version
 
-PYDANTIC_VERSION = semver.parse_version_info(pydantic.__version__)
+PYDANTIC_VERSION = Version(pydantic.__version__)
 try:
     from pydantic_core import CoreSchema, core_schema
 except ImportError:
-    if PYDANTIC_VERSION >= (2,):
+    if PYDANTIC_VERSION.major >= 2:
         raise
 
 if TYPE_CHECKING:
@@ -144,7 +144,7 @@ def Vector(
                 raise TypeError("A list of numbers or numpy.ndarray is needed")
             return cls(v)
 
-        if PYDANTIC_VERSION < (2, 0):
+        if PYDANTIC_VERSION.major < 2:
 
             @classmethod
             def __modify_schema__(cls, field_schema: Dict[str, Any]):

python/python/lancedb/query.py

@@ -1113,11 +1113,22 @@ class AsyncQueryBase(object):
         self._inner.limit(limit)
         return self
 
-    async def to_batches(self) -> AsyncRecordBatchReader:
+    async def to_batches(
+        self, *, max_batch_length: Optional[int] = None
+    ) -> AsyncRecordBatchReader:
         """
         Execute the query and return the results as an Apache Arrow RecordBatchReader.
+
+        Parameters
+        ----------
+
+        max_batch_length: Optional[int]
+            The maximum number of selected records in a single RecordBatch object.
+            If not specified, a default batch length is used.
+            It is possible for batches to be smaller than the provided length if the
+            underlying data is stored in smaller chunks.
         """
-        return AsyncRecordBatchReader(await self._inner.execute())
+        return AsyncRecordBatchReader(await self._inner.execute(max_batch_length))
 
     async def to_arrow(self) -> pa.Table:
         """

python/python/lancedb/rerankers/cohere.py

@@ -1,5 +1,5 @@
 import os
-import semver
+from packaging.version import Version
 from functools import cached_property
 from typing import Union
 
@@ -44,9 +44,8 @@ class CohereReranker(Reranker):
     def _client(self):
         cohere = attempt_import_or_raise("cohere")
         # ensure version is at least 0.5.0
-        if (
-            hasattr(cohere, "__version__")
-            and semver.compare(cohere.__version__, "5.0.0") < 0
+        if hasattr(cohere, "__version__") and Version(cohere.__version__) < Version(
+            "0.5.0"
         ):
             raise ValueError(
                 f"cohere version must be at least 0.5.0, found {cohere.__version__}"

python/python/lancedb/table.py

@@ -58,7 +58,7 @@ if TYPE_CHECKING:
     import PIL
     from lance.dataset import CleanupStats, ReaderLike
 
-    from ._lancedb import Table as LanceDBTable
+    from ._lancedb import Table as LanceDBTable, OptimizeStats
     from .db import LanceDBConnection
     from .index import BTree, IndexConfig, IvfPq
 
@@ -2377,6 +2377,49 @@ class AsyncTable:
         """
         await self._inner.restore()
 
+    async def optimize(
+        self, *, cleanup_older_than: Optional[timedelta] = None
+    ) -> OptimizeStats:
+        """
+        Optimize the on-disk data and indices for better performance.
+
+        Modeled after ``VACUUM`` in PostgreSQL.
+
+        Optimization covers three operations:
+
+         * Compaction: Merges small files into larger ones
+         * Prune: Removes old versions of the dataset
+         * Index: Optimizes the indices, adding new data to existing indices
+
+        Parameters
+        ----------
+        cleanup_older_than: timedelta, optional default 7 days
+            All files belonging to versions older than this will be removed.  Set
+            to 0 days to remove all versions except the latest.  The latest version
+            is never removed.
+
+        Experimental API
+        ----------------
+
+        The optimization process is undergoing active development and may change.
+        Our goal with these changes is to improve the performance of optimization and
+        reduce the complexity.
+
+        That being said, it is essential today to run optimize if you want the best
+        performance.  It should be stable and safe to use in production, but it our
+        hope that the API may be simplified (or not even need to be called) in the
+        future.
+
+        The frequency an application shoudl call optimize is based on the frequency of
+        data modifications.  If data is frequently added, deleted, or updated then
+        optimize should be run frequently.  A good rule of thumb is to run optimize if
+        you have added or modified 100,000 or more records or run more than 20 data
+        modification operations.
+        """
+        if cleanup_older_than is not None:
+            cleanup_older_than = round(cleanup_older_than.total_seconds() * 1000)
+        return await self._inner.optimize(cleanup_older_than)
+
     async def list_indices(self) -> IndexConfig:
         """
         List all indices that have been created with Self::create_index