mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-26 07:58:31 +00:00
Compare commits
102 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 2f8b3bccd2 | |||
| 31b68dd21b | |||
| 60ba79e114 | |||
| 83c84b3b6c | |||
| 3a3e488b34 | |||
| be290447d9 | |||
| 79ba076429 | |||
| 42dfecc582 | |||
| ec21e37040 | |||
| 6ba80a960c | |||
| 11f24b1df4 | |||
| 2ba7407dc3 | |||
| 607e556927 | |||
| 564e5d0d56 | |||
| dd5cb4d805 | |||
| dbc3687c7b | |||
| ec80acb668 | |||
| fc44535cee | |||
| 4048150fdd | |||
| 2922c171f7 | |||
| c5f9efefe9 | |||
| f4c668e244 | |||
| b1cfe6edb1 | |||
| 33b0e2dacd | |||
| 6c4b67b9ef | |||
| 1990e6a1ca | |||
| 938beb35e9 | |||
| 6a8146f376 | |||
| 02fbe212e8 | |||
| 6a9553a902 | |||
| 79b65247bd | |||
| 7db2628c87 | |||
| bf35358be6 | |||
| a4d34d7b5d | |||
| 261e9272f6 | |||
| 001237c7a4 | |||
| 369b10a377 | |||
| 1c3cd1d918 | |||
| 9707966943 | |||
| 62fe413a52 | |||
| 1493ece3de | |||
| e6444ecc05 | |||
| cc0139c136 | |||
| b20696ef9c | |||
| 772bdeced8 | |||
| c1a3fa7f51 | |||
| 0ba82873c5 | |||
| 3af51541a0 | |||
| 2c06a48bd8 | |||
| ac8b28c010 | |||
| 173f889d2a | |||
| 03b52e5877 | |||
| 798e5364fb | |||
| f1f34dfdd3 | |||
| 123c921c4f | |||
| 99a68db78c | |||
| 9e73d440a3 | |||
| 3956d9dbfa | |||
| 16e1967efc | |||
| 27dd92c67e | |||
| 9e2e711c7a | |||
| c3176a47ce | |||
| bae68d45ca | |||
| 7357d63e87 | |||
| 624a75edf7 | |||
| c7ea91f3ea | |||
| 8e24dd3828 | |||
| f79dc017c4 | |||
| e6ae93f52a | |||
| 3dd9c598e9 | |||
| 9e26bf3fba | |||
| 93354baf34 | |||
| 05602ec7d5 | |||
| e3b472c212 | |||
| a6418b6cb9 | |||
| dd2b11eda2 | |||
| 5a1015ba72 | |||
| 48945d0658 | |||
| 77208fd464 | |||
| b505dc1315 | |||
| 7dfdfe6401 | |||
| 4dc2d9a0f2 | |||
| 1ad6ce3a4e | |||
| 03b26d585b | |||
| f7feed48c3 | |||
| e5f489818b | |||
| 98a52267a2 | |||
| ff50e698cf | |||
| b799ebaa69 | |||
| 72fc660f9e | |||
| 1ebde1f06c | |||
| 29c030f865 | |||
| ff6ff09998 | |||
| 119b9baf90 | |||
| ba4558a64f | |||
| f655f62e09 | |||
| bf15655c83 | |||
| a00edef0e6 | |||
| 1b2670443e | |||
| 9dc5ec03aa | |||
| 18760f74cd | |||
| c9d07ef6fc |
+8
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.32.0-beta.2"
|
||||
current_version = "0.37.1-beta.0"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
@@ -75,6 +75,13 @@ filename = "nodejs/Cargo.toml"
|
||||
replace = "\nversion = \"{new_version}\""
|
||||
search = "\nversion = \"{current_version}\""
|
||||
|
||||
# The Python package takes its version from here (pyproject.toml declares
|
||||
# `dynamic = ["version"]`, so maturin reads it out of the crate manifest).
|
||||
[[tool.bumpversion.files]]
|
||||
filename = "python/Cargo.toml"
|
||||
replace = "\nversion = \"{new_version}\""
|
||||
search = "\nversion = \"{current_version}\""
|
||||
|
||||
# Java documentation
|
||||
[[tool.bumpversion.files]]
|
||||
filename = "docs/src/java/java.md"
|
||||
|
||||
@@ -27,19 +27,31 @@ runs:
|
||||
# Extract failed job names
|
||||
FAILED_JOBS=$(echo "$JOB_RESULTS" | jq -r 'to_entries | map(select(.value.result == "failure")) | map(.key) | join(", ")')
|
||||
|
||||
# Create issue with workflow name, failed jobs, and run URL
|
||||
gh issue create \
|
||||
--title "$WORKFLOW_NAME Failed ($FAILED_JOBS)" \
|
||||
--body "The workflow **$WORKFLOW_NAME** failed during execution.
|
||||
TITLE="$WORKFLOW_NAME Failed ($FAILED_JOBS)"
|
||||
|
||||
# This action now also runs on nightly schedules, so a breakage that
|
||||
# persists for a few days would otherwise file one issue per night.
|
||||
# Comment on the open report instead when one already exists.
|
||||
EXISTING=$(gh issue list --state open --label ci --limit 100 --json number,title \
|
||||
| jq -r --arg title "$TITLE" 'map(select(.title == $title)) | .[0].number // empty')
|
||||
|
||||
if [ -n "$EXISTING" ]; then
|
||||
gh issue comment "$EXISTING" --body "Failed again: $RUN_URL"
|
||||
echo "Commented on existing issue #$EXISTING"
|
||||
else
|
||||
gh issue create \
|
||||
--title "$TITLE" \
|
||||
--body "The workflow **$WORKFLOW_NAME** failed during execution.
|
||||
|
||||
**Failed jobs:** $FAILED_JOBS
|
||||
|
||||
**Run URL:** $RUN_URL
|
||||
|
||||
Please investigate the failed jobs and address any issues." \
|
||||
--label "ci"
|
||||
--label "ci"
|
||||
|
||||
echo "Issue created successfully"
|
||||
echo "Issue created successfully"
|
||||
fi
|
||||
else
|
||||
echo "No job failures detected, skipping issue creation"
|
||||
fi
|
||||
|
||||
@@ -6,7 +6,6 @@ on:
|
||||
# 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
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
name: Check doc links
|
||||
|
||||
# Checking external links is inherently noisy: third-party sites rate-limit
|
||||
# automated clients, reject non-browser user agents, and go down temporarily.
|
||||
# Blocking pull requests on that trades a lot of false failures for very little
|
||||
# signal, so this runs on a schedule and reports findings in a single tracking
|
||||
# issue instead of failing anyone's build.
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 7 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
# The report lives in one repository-global issue, so runs must not overlap: a
|
||||
# lookup racing a create produces duplicate issues, and a healthy run closing
|
||||
# the issue while a failing run only rewrites its body would leave a broken
|
||||
# report closed. The group is deliberately ref-independent so that a manual
|
||||
# dispatch serializes against the scheduled run.
|
||||
concurrency:
|
||||
group: docs-link-check
|
||||
cancel-in-progress: false
|
||||
|
||||
permissions: {}
|
||||
|
||||
env:
|
||||
REPORT_TITLE: "Docs link checker report"
|
||||
|
||||
jobs:
|
||||
scan:
|
||||
name: Scan links
|
||||
runs-on: ubuntu-24.04
|
||||
# lychee-action is pinned by SHA, but its wrapper downloads the lychee
|
||||
# release tarball at run time without verifying a digest, and hands the
|
||||
# resulting binary a GitHub token. Release assets remain replaceable, so
|
||||
# that binary is confined to a job whose token can only read public
|
||||
# content; everything that writes runs in the report job below.
|
||||
permissions:
|
||||
contents: read
|
||||
outputs:
|
||||
exit_code: ${{ steps.lychee.outputs.exit_code }}
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
# workflow_dispatch can run from any ref, but the report is
|
||||
# repository-global. Always measure the default branch so a manual
|
||||
# run from a topic branch cannot close a report that main warrants,
|
||||
# or overwrite it with branch-only findings.
|
||||
ref: ${{ github.event.repository.default_branch }}
|
||||
persist-credentials: false
|
||||
|
||||
- name: Check links
|
||||
id: lychee
|
||||
uses: lycheeverse/lychee-action@e7477775783ea5526144ba13e8db5eec57747ce8 # v2.9.0
|
||||
with:
|
||||
# Restricted to http(s) on purpose. Much of docs/src is generated
|
||||
# API reference (the js/ tree comes from `npm run docs` in nodejs)
|
||||
# and the hand-written pages use mkdocstrings cross-references and
|
||||
# nav-relative paths that only resolve in the site mkdocs builds,
|
||||
# not in this checkout, so relative links would be reported as
|
||||
# broken on every run.
|
||||
args: >-
|
||||
--scheme https
|
||||
--scheme http
|
||||
--no-progress
|
||||
--max-retries 3
|
||||
--timeout 20
|
||||
'docs/src/**/*.md'
|
||||
format: json
|
||||
output: ./lychee/out.json
|
||||
jobSummary: false
|
||||
# The report, not a red build, is the signal for broken links. The
|
||||
# validation step below still fails the run if the check itself
|
||||
# breaks.
|
||||
fail: false
|
||||
|
||||
- name: Validate report
|
||||
# lychee does not reserve exit code 2 for broken links: its CLI
|
||||
# parser also exits 2 on an invalid option, before any link was
|
||||
# checked or any report written. Only a parseable report whose
|
||||
# counts agree with the exit code counts as a link verdict; anything
|
||||
# else fails here, and the report job below is skipped entirely, so
|
||||
# the tracking issue is never touched. Exit 2 covers timeouts as
|
||||
# well as errors, and a timed-out host is exactly the transient
|
||||
# unavailability this report exists to surface, so both count as
|
||||
# findings. Requiring total > 0 also catches a glob that silently
|
||||
# stopped matching any file.
|
||||
if: steps.lychee.outputs.exit_code == 0 || steps.lychee.outputs.exit_code == 2
|
||||
env:
|
||||
EXIT_CODE: ${{ steps.lychee.outputs.exit_code }}
|
||||
run: |
|
||||
jq -e --argjson code "$EXIT_CODE" '
|
||||
(.total > 0) and
|
||||
(if $code == 0
|
||||
then .errors == 0 and .timeouts == 0
|
||||
and (.error_map | length == 0) and (.timeout_map | length == 0)
|
||||
else (.errors + .timeouts) > 0
|
||||
and ((.error_map | length) + (.timeout_map | length)) > 0
|
||||
end)
|
||||
' ./lychee/out.json
|
||||
|
||||
- name: Upload report
|
||||
if: steps.lychee.outputs.exit_code == 2
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: link-report
|
||||
path: ./lychee/out.json
|
||||
retention-days: 7
|
||||
|
||||
report:
|
||||
name: Update report issue
|
||||
needs: scan
|
||||
runs-on: ubuntu-24.04
|
||||
# Deliberately no checkout: this job needs the report artifact and the
|
||||
# issues API, not the repository contents.
|
||||
permissions:
|
||||
issues: write
|
||||
env:
|
||||
EXIT_CODE: ${{ needs.scan.outputs.exit_code }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
steps:
|
||||
- name: Classify checker result
|
||||
# lychee exits 0 when every link resolves and 2 when links fail,
|
||||
# both already cross-checked against the report by the scan job's
|
||||
# validation step. Anything else (1 runtime, 3 bad config) means the
|
||||
# check never produced a link verdict, which must surface as a failed
|
||||
# run rather than be published as "broken documentation links".
|
||||
run: |
|
||||
case "$EXIT_CODE" in
|
||||
0|2)
|
||||
echo "lychee exit code $EXIT_CODE"
|
||||
;;
|
||||
*)
|
||||
echo "::error::lychee exited with '$EXIT_CODE': the link check did not complete. Leaving the report issue untouched."
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
- name: Find existing report issue
|
||||
id: report
|
||||
# Matched on title alone, and through search rather than a listing:
|
||||
# the issue action applies labels in a separate call after creating the
|
||||
# issue, so a label filter misses a half-created report, and this
|
||||
# repository has far more open issues than one listing page holds.
|
||||
# Closed issues are included because a healthy run closes the report:
|
||||
# an open-only lookup would forget that identity and the next failing
|
||||
# run would open a duplicate. The oldest match stays the canonical
|
||||
# report and is reopened below when links break again.
|
||||
run: |
|
||||
match=$(gh issue list --repo "$GITHUB_REPOSITORY" --state all \
|
||||
--search "in:title \"$REPORT_TITLE\" author:app/github-actions" \
|
||||
--limit 50 --json number,title,state \
|
||||
--jq "[.[] | select(.title == \"$REPORT_TITLE\")] | sort_by(.number) | first // empty")
|
||||
echo "number=$(jq -r '.number // empty' <<<"$match")" >> "$GITHUB_OUTPUT"
|
||||
echo "state=$(jq -r '.state // empty' <<<"$match")" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Download report
|
||||
if: env.EXIT_CODE == 2
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: link-report
|
||||
path: ./lychee
|
||||
|
||||
- name: Compose report
|
||||
if: env.EXIT_CODE == 2
|
||||
run: |
|
||||
run_url="$GITHUB_SERVER_URL/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID"
|
||||
{
|
||||
echo "Broken documentation links found by [\`$GITHUB_WORKFLOW\`]($run_url)."
|
||||
echo
|
||||
echo "This issue is rewritten by every scheduled run and closed automatically once all links resolve."
|
||||
echo
|
||||
echo "Entries can be false positives: some sites rate-limit or block automated clients while working fine in a browser. Confirm before editing the docs, and add persistent offenders to \`--exclude\` in \`.github/workflows/docs-link-check.yml\`."
|
||||
echo
|
||||
# Timeouts are reported alongside errors: entries land in
|
||||
# timeout_map with a status text instead of an HTTP code.
|
||||
jq -r '
|
||||
"\(.errors) of \(.total) links failed, \(.timeouts) timed out.",
|
||||
"",
|
||||
([(.error_map | to_entries[]), (.timeout_map | to_entries[])]
|
||||
| group_by(.key)[] |
|
||||
"### Errors in \(.[0].key)",
|
||||
"",
|
||||
(map(.value[])[] | "* [\(.status.code // .status.text // "ERR")] <\(.url)> — \(.status.details // .status.text // "unknown error")"),
|
||||
"")
|
||||
' ./lychee/out.json
|
||||
} > ./lychee/issue.md
|
||||
|
||||
- name: Reopen report issue
|
||||
# A healthy run closes the report, and the issue action below only
|
||||
# rewrites the body of whatever number it is given. Without an
|
||||
# explicit reopen, the 2 -> 0 -> 2 sequence would keep rewriting a
|
||||
# closed issue while links are broken. A CLOSED state implies the
|
||||
# lookup found a canonical issue, so no separate emptiness check.
|
||||
if: env.EXIT_CODE == 2 && steps.report.outputs.state == 'CLOSED'
|
||||
env:
|
||||
ISSUE_NUMBER: ${{ steps.report.outputs.number }}
|
||||
run: |
|
||||
run_url="$GITHUB_SERVER_URL/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID"
|
||||
gh issue reopen "$ISSUE_NUMBER" --repo "$GITHUB_REPOSITORY" \
|
||||
--comment "Broken documentation links found again in [the latest run]($run_url)."
|
||||
|
||||
- name: Report broken links
|
||||
if: env.EXIT_CODE == 2
|
||||
uses: peter-evans/create-issue-from-file@fca9117c27cdc29c6c4db3b86c48e4115a786710 # v6.0.0
|
||||
with:
|
||||
# Empty on the first failing run, which creates the issue; afterwards
|
||||
# the same issue is updated in place.
|
||||
issue-number: ${{ steps.report.outputs.number }}
|
||||
title: ${{ env.REPORT_TITLE }}
|
||||
content-filepath: ./lychee/issue.md
|
||||
labels: documentation
|
||||
|
||||
- name: Close report issue once links are healthy
|
||||
# An OPEN state implies the lookup found a canonical issue; a report
|
||||
# that is already closed needs nothing.
|
||||
if: env.EXIT_CODE == 0 && steps.report.outputs.state == 'OPEN'
|
||||
env:
|
||||
ISSUE_NUMBER: ${{ steps.report.outputs.number }}
|
||||
run: |
|
||||
run_url="$GITHUB_SERVER_URL/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID"
|
||||
gh issue close "$ISSUE_NUMBER" --repo "$GITHUB_REPOSITORY" \
|
||||
--comment "All documentation links resolved in [the latest run]($run_url)."
|
||||
@@ -0,0 +1,85 @@
|
||||
name: GitHub Release
|
||||
|
||||
# All SDKs share one version, so a single `vX.Y.Z` tag produces a single GitHub
|
||||
# release covering all of them. The per-package publish workflows (PyPI, NPM,
|
||||
# Cargo, Maven) trigger off the same tag independently.
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- "v*"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
gh-release:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
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"
|
||||
echo "prerelease=true" >> $GITHUB_OUTPUT
|
||||
|
||||
# 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"
|
||||
echo "prerelease=false" >> $GITHUB_OUTPUT
|
||||
# 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:
|
||||
# Marking betas as pre-releases keeps them from taking the "Latest"
|
||||
# badge on the releases page.
|
||||
prerelease: ${{ steps.extract_version.outputs.prerelease }}
|
||||
make_latest: ${{ steps.extract_version.outputs.prerelease == 'false' }}
|
||||
tag_name: ${{ steps.extract_version.outputs.tag }}
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
generate_release_notes: false
|
||||
name: LanceDB v${{ steps.extract_version.outputs.version }}
|
||||
body: ${{ steps.release_notes.outputs.changelog }}
|
||||
@@ -1,13 +1,14 @@
|
||||
name: Create release commit
|
||||
|
||||
# This workflow increments versions, tags the version, and pushes it.
|
||||
# This workflow increments the version, tags it, and pushes it. All SDKs share
|
||||
# a single version, so one tag releases all of them.
|
||||
# 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.
|
||||
# breaking changes since the last minor increment. A breaking change in any SDK
|
||||
# bumps the minor version for all of them. If you wish to bypass this check, you
|
||||
# can manually increment the version and push the tag.
|
||||
on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
@@ -24,16 +25,6 @@ on:
|
||||
options:
|
||||
- preview
|
||||
- stable
|
||||
python:
|
||||
description: 'Make a Python release'
|
||||
required: true
|
||||
default: true
|
||||
type: boolean
|
||||
other:
|
||||
description: 'Make a Node/Rust/Java release'
|
||||
required: true
|
||||
default: true
|
||||
type: boolean
|
||||
bump-minor:
|
||||
description: 'Bump minor version'
|
||||
required: true
|
||||
@@ -65,25 +56,12 @@ jobs:
|
||||
run: |
|
||||
git config user.name 'Lance Release'
|
||||
git config user.email 'lance-dev@lancedb.com'
|
||||
- name: Bump Python version
|
||||
if: ${{ inputs.python }}
|
||||
working-directory: python
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
# 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 }}
|
||||
- name: Bump version
|
||||
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
|
||||
bash ci/bump_version.sh ${{ inputs.type }} ${{ inputs.bump-minor }}
|
||||
bash ci/update_lockfiles.sh --amend
|
||||
- name: Push new version tag
|
||||
if: ${{ !inputs.dry_run }}
|
||||
|
||||
@@ -61,6 +61,11 @@ jobs:
|
||||
sudo apt update
|
||||
sudo apt install -y protobuf-compiler libssl-dev
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Format Rust
|
||||
run: cargo fmt --all -- --check
|
||||
- name: Lint Rust
|
||||
@@ -103,6 +108,11 @@ jobs:
|
||||
cache: 'pnpm'
|
||||
cache-dependency-path: nodejs/pnpm-lock.yaml
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
sudo apt update
|
||||
@@ -182,6 +192,11 @@ jobs:
|
||||
cache-dependency-path: nodejs/pnpm-lock.yaml
|
||||
- uses: dtolnay/rust-toolchain@stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
brew install protobuf
|
||||
|
||||
@@ -10,10 +10,16 @@ permissions:
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
tags:
|
||||
- "v*"
|
||||
# The cross-compiled targets (musl especially) break from toolchain and
|
||||
# dependency changes that nothing else in CI catches, and discovering that
|
||||
# mid-release is expensive. A nightly run keeps that signal while dropping
|
||||
# the full 8-target release matrix from all ~90 pushes to main each month.
|
||||
# `report-failure` files an issue when a nightly breaks.
|
||||
schedule:
|
||||
- cron: "0 8 * * *"
|
||||
workflow_dispatch:
|
||||
pull_request:
|
||||
# This should trigger a dry run (we skip the final publish step)
|
||||
paths:
|
||||
@@ -26,73 +32,6 @@ concurrency:
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
gh-release:
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
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 }}
|
||||
|
||||
build-lancedb:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
@@ -101,9 +40,18 @@ jobs:
|
||||
- target: aarch64-apple-darwin
|
||||
host: macos-latest
|
||||
features: fp16kernels
|
||||
pre_build: brew install protobuf
|
||||
pre_build: |-
|
||||
brew install protobuf
|
||||
# Fat LTO (the workspace default in .cargo/config.toml) is
|
||||
# single-threaded and is the peak-memory step of the build. On
|
||||
# this runner it accounted for ~111 of the job's ~113 minutes,
|
||||
# making it the critical path of the entire publish pipeline.
|
||||
# ThinLTO parallelizes it across the runner's cores, for a few
|
||||
# percent of runtime performance.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: x86_64-pc-windows-msvc
|
||||
host: windows-2025-8x-x64
|
||||
host: windows-2025
|
||||
features: ","
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc ninja nasm
|
||||
@@ -111,19 +59,19 @@ jobs:
|
||||
# There is an issue where choco doesn't add nasm to the path
|
||||
export PATH="$PATH:/c/Program Files/NASM"
|
||||
nasm -v
|
||||
# Fat LTO of the cdylib is single-threaded and the peak-memory
|
||||
# step of the build, and had started hitting rustc-LLVM OOM on the
|
||||
# Windows runners. ThinLTO parallelizes it across the runner's
|
||||
# cores and keeps peak memory well under the limit.
|
||||
# See the ThinLTO note on aarch64-apple-darwin above. Keeping
|
||||
# peak memory down is also what lets this run on the standard
|
||||
# 4-core runner: the 8-core larger runner was only needed to
|
||||
# stop fat LTO from OOMing rustc-LLVM.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: aarch64-pc-windows-msvc
|
||||
host: windows-2025-8x-x64
|
||||
host: windows-2025
|
||||
features: ","
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc
|
||||
rustup target add aarch64-pc-windows-msvc
|
||||
# See ThinLTO note on the x86_64-pc-windows-msvc target above.
|
||||
# See the ThinLTO note on aarch64-apple-darwin above.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: x86_64-unknown-linux-gnu
|
||||
@@ -198,16 +146,49 @@ jobs:
|
||||
with:
|
||||
toolchain: stable
|
||||
targets: ${{ matrix.settings.target }}
|
||||
- name: Cache cargo
|
||||
uses: actions/cache@v5
|
||||
# These builds were entirely uncached: the old key was static, so
|
||||
# `actions/cache` (which only writes on a miss) could never refresh it,
|
||||
# and the multi-GB whole-`target/` copy it tried to store never fit the
|
||||
# repo's cache budget, so no entry was ever saved. rust-cache prunes
|
||||
# `target/` to dependency artifacts and keys on Cargo.lock plus the rustc
|
||||
# version, which both fixes the key and keeps entries a sane size.
|
||||
#
|
||||
# This caches dependency *compilation* only. The LTO link of the cdylib
|
||||
# re-runs regardless, since the local crate changes every time, so the
|
||||
# win is larger on the non-LTO jobs than here.
|
||||
- name: Cache cargo (native builds)
|
||||
uses: Swatinem/rust-cache@v2
|
||||
if: ${{ !matrix.settings.docker }}
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry/index/
|
||||
~/.cargo/registry/cache/
|
||||
~/.cargo/git/db/
|
||||
.cargo-cache
|
||||
target/
|
||||
key: nodejs-${{ matrix.settings.target }}-cargo-${{ matrix.settings.host }}
|
||||
# The release profile and per-target dirs differ from what the test
|
||||
# workflows cache, so these need to be separate entries.
|
||||
key: release-${{ matrix.settings.target }}
|
||||
# Only the nightly run on main writes, so tag and PR runs restore a
|
||||
# warm entry without every dependabot PR writing its own (which would
|
||||
# be unreadable elsewhere anyway, since GitHub scopes caches to the
|
||||
# creating ref). The nightly cadence also keeps entries inside
|
||||
# GitHub's 7-day eviction window, which a tag-only trigger would not.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
# Docker builds can use rust-cache too. `target/` already lives on the
|
||||
# host because the whole workspace is bind-mounted into the container, and
|
||||
# rust-cache's prune and save run host-side, so they can manage it -- which
|
||||
# is what keeps the entry to dependency artifacts rather than a multi-GB
|
||||
# copy of everything.
|
||||
#
|
||||
# Two differences from the native builds. The container's CARGO_HOME is
|
||||
# bind-mounted from `.cargo-cache` rather than the host's ~/.cargo, so that
|
||||
# has to be cached explicitly. And the key is derived from the *host* rustc
|
||||
# version, which is not the compiler that produced these artifacts; that is
|
||||
# safe because cargo fingerprints the real compiler and rebuilds on a
|
||||
# mismatch, it just means a base-image toolchain bump costs one cold build
|
||||
# instead of invalidating the key.
|
||||
- name: Cache cargo (docker builds)
|
||||
uses: Swatinem/rust-cache@v2
|
||||
if: ${{ matrix.settings.docker }}
|
||||
with:
|
||||
key: docker-${{ matrix.settings.target }}
|
||||
cache-directories: .cargo-cache
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- name: Install Zig
|
||||
@@ -225,9 +206,13 @@ jobs:
|
||||
if: ${{ matrix.settings.docker }}
|
||||
with:
|
||||
image: ${{ matrix.settings.docker }}
|
||||
# All three mounts must live under `.cargo-cache`, which is what the
|
||||
# cache step above saves. Previously the registry mounts pointed at
|
||||
# `.cargo/...`, a path nothing cached, so the container re-downloaded
|
||||
# the whole crate registry on every run.
|
||||
options: "--user 0:0 -v ${{ github.workspace }}/.cargo-cache/git/db:/usr/local/cargo/git/db \
|
||||
-v ${{ github.workspace }}/.cargo/registry/cache:/usr/local/cargo/registry/cache \
|
||||
-v ${{ github.workspace }}/.cargo/registry/index:/usr/local/cargo/registry/index \
|
||||
-v ${{ github.workspace }}/.cargo-cache/registry/cache:/usr/local/cargo/registry/cache \
|
||||
-v ${{ github.workspace }}/.cargo-cache/registry/index:/usr/local/cargo/registry/index \
|
||||
-v ${{ github.workspace }}:/build -w /build/nodejs"
|
||||
run: |
|
||||
set -e
|
||||
@@ -239,6 +224,16 @@ jobs:
|
||||
--js ../lancedb/native.js \
|
||||
--strip \
|
||||
--output-dir dist/
|
||||
# The container runs as root (`--user 0:0`), so everything it wrote to the
|
||||
# mounted cache dirs is root-owned. rust-cache's post step runs as the
|
||||
# runner user and has to both read these and delete from them while
|
||||
# pruning, so hand them back before it runs.
|
||||
- name: Take ownership of docker build output
|
||||
if: ${{ matrix.settings.docker }}
|
||||
run: |
|
||||
sudo chown -R "$(id -u):$(id -g)" \
|
||||
"${{ github.workspace }}/.cargo-cache" \
|
||||
"${{ github.workspace }}/target"
|
||||
- name: Build
|
||||
run: |
|
||||
${{ matrix.settings.pre_build }}
|
||||
@@ -252,6 +247,15 @@ jobs:
|
||||
--output-dir dist/
|
||||
if: ${{ !matrix.settings.docker }}
|
||||
shell: bash
|
||||
# The standard Windows runners have ~14 GB free, and a release `target/`
|
||||
# for this workspace is a large fraction of that. Report the remaining
|
||||
# headroom so a build that only just fits is visible before a dependency
|
||||
# bump turns it into a failed release. `always()` so the numbers are
|
||||
# still there when the build is what ran out of space.
|
||||
- name: Report disk headroom
|
||||
if: always()
|
||||
run: df -h
|
||||
shell: bash
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
@@ -402,7 +406,9 @@ jobs:
|
||||
name: Report Workflow Failure
|
||||
runs-on: ubuntu-latest
|
||||
needs: [build-lancedb, test-lancedb, publish]
|
||||
if: always() && failure() && startsWith(github.ref, 'refs/tags/v')
|
||||
# Nightly runs are the only thing watching the cross-compiled targets now,
|
||||
# so they have to report failures too or the signal is silently lost.
|
||||
if: always() && failure() && (startsWith(github.ref, 'refs/tags/v') || github.event_name == 'schedule')
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
|
||||
@@ -3,7 +3,7 @@ name: PyPI Publish
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'python-v*'
|
||||
- 'v*'
|
||||
pull_request:
|
||||
# This should trigger a dry run (we skip the final publish step)
|
||||
paths:
|
||||
@@ -20,6 +20,12 @@ env:
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
# Without this, a force-push to a PR leaves the previous run going -- including
|
||||
# a ~74 minute Windows job and a billed arm64 wheel build.
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
linux:
|
||||
name: Python ${{ matrix.config.package_name }} ${{ matrix.config.platform }} manylinux${{ matrix.config.manylinux }}
|
||||
@@ -72,7 +78,7 @@ jobs:
|
||||
package-name: ${{ matrix.config.package_name }}
|
||||
rustflags: ${{ matrix.config.rustflags }}
|
||||
- uses: actions/upload-artifact@v7
|
||||
if: startsWith(github.ref, 'refs/tags/python-v')
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
with:
|
||||
name: wheels-linux-${{ matrix.config.package_name }}-${{ matrix.config.platform }}-${{ matrix.config.manylinux }}
|
||||
path: target/wheels/*.whl
|
||||
@@ -101,7 +107,7 @@ jobs:
|
||||
python-minor-version: 10
|
||||
args: "--release --strip --target ${{ matrix.config.target }} --features fp16kernels"
|
||||
- uses: actions/upload-artifact@v7
|
||||
if: startsWith(github.ref, 'refs/tags/python-v')
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
with:
|
||||
name: wheels-mac-${{ matrix.config.target }}
|
||||
path: target/wheels/lancedb-*.whl
|
||||
@@ -122,19 +128,26 @@ jobs:
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.13"
|
||||
# NOTE: caching cargo here would be a no-op. This workflow only runs on
|
||||
# tags and PRs, and GitHub only lets a run restore caches from its own ref
|
||||
# or the default branch -- so with no run on main there is nothing that
|
||||
# can populate an entry the release build would be allowed to read. Fixing
|
||||
# this needs a main/nightly trigger (which would also catch wheel-build
|
||||
# breakage before a release); the ~74 minutes here is otherwise dominated
|
||||
# by the fat-LTO link, which no cache avoids.
|
||||
- uses: ./.github/workflows/build_windows_wheel
|
||||
with:
|
||||
python-minor-version: 10
|
||||
args: "--release --strip"
|
||||
- uses: actions/upload-artifact@v7
|
||||
if: startsWith(github.ref, 'refs/tags/python-v')
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
with:
|
||||
name: wheels-windows
|
||||
path: target/wheels/lancedb-*.whl
|
||||
if-no-files-found: error
|
||||
publish:
|
||||
name: Publish wheels
|
||||
if: startsWith(github.ref, 'refs/tags/python-v')
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
needs: [linux, mac, windows]
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
@@ -183,72 +196,6 @@ jobs:
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: target/wheels/
|
||||
gh-release:
|
||||
if: startsWith(github.ref, 'refs/tags/python-v')
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
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 }}
|
||||
report-failure:
|
||||
name: Report Workflow Failure
|
||||
runs-on: ubuntu-latest
|
||||
@@ -256,7 +203,7 @@ jobs:
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
if: always() && failure() && startsWith(github.ref, 'refs/tags/python-v')
|
||||
if: always() && failure() && startsWith(github.ref, 'refs/tags/v')
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
|
||||
@@ -108,6 +108,15 @@ jobs:
|
||||
run: |
|
||||
sudo apt update
|
||||
sudo apt install -y protobuf-compiler
|
||||
# `pip install -e .` builds the extension with maturin, which is most of
|
||||
# this job's ~33 minutes. It had no Rust cache, so every dependency was
|
||||
# recompiled from scratch on every run.
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install
|
||||
run: |
|
||||
pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -e .[tests,dev,embeddings]
|
||||
@@ -168,6 +177,14 @@ jobs:
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.13"
|
||||
# maturin runs cargo natively on macOS (docker is Linux-only), so the host
|
||||
# target dir is cacheable. This job had no Rust cache.
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- uses: ./.github/workflows/build_mac_wheel
|
||||
with:
|
||||
args: --profile ci
|
||||
@@ -197,6 +214,14 @@ jobs:
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.13"
|
||||
# maturin runs cargo natively on Windows (docker is Linux-only), so the
|
||||
# host target dir is cacheable. This job had no Rust cache at all and so
|
||||
# rebuilt every dependency from scratch on every run.
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. The repo sits at
|
||||
# GitHub's cache cap, so per-PR saves just evict main's entries.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- uses: ./.github/workflows/build_windows_wheel
|
||||
with:
|
||||
args: --profile ci
|
||||
@@ -224,6 +249,14 @@ jobs:
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.10"
|
||||
# As with Doctest, `pip install -e .` compiles the extension and this job
|
||||
# had no Rust cache, which is most of its ~37 minutes.
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install lancedb
|
||||
run: |
|
||||
pip install "pydantic<2"
|
||||
|
||||
+53
-13
@@ -48,6 +48,11 @@ jobs:
|
||||
with:
|
||||
components: rustfmt, clippy
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
sudo apt update
|
||||
@@ -89,6 +94,11 @@ jobs:
|
||||
run: rm -f Cargo.lock
|
||||
- uses: rui314/setup-mold@v1
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
sudo apt update
|
||||
@@ -118,6 +128,11 @@ jobs:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
sudo apt update
|
||||
@@ -175,6 +190,11 @@ jobs:
|
||||
- name: CPU features
|
||||
run: sysctl -a | grep cpu
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install dependencies
|
||||
run: brew install protobuf
|
||||
- name: Run tests
|
||||
@@ -187,12 +207,19 @@ jobs:
|
||||
cargo test --profile ci --features $ALL_FEATURES --locked
|
||||
|
||||
windows:
|
||||
runs-on: windows-2022
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
target:
|
||||
- x86_64-pc-windows-msvc
|
||||
- aarch64-pc-windows-msvc
|
||||
include:
|
||||
- target: x86_64-pc-windows-msvc
|
||||
runner: windows-2022
|
||||
# windows-11-arm is a standard runner, so it is free on public repos.
|
||||
# Running natively lets the aarch64 tests actually execute -- this
|
||||
# job used to cross-compile them and then skip the test step, paying
|
||||
# full codegen and link cost for a compile check.
|
||||
- target: aarch64-pc-windows-msvc
|
||||
runner: windows-11-arm
|
||||
runs-on: ${{ matrix.runner }}
|
||||
defaults:
|
||||
run:
|
||||
working-directory: rust/lancedb
|
||||
@@ -201,6 +228,11 @@ jobs:
|
||||
- name: Set target
|
||||
run: rustup target add ${{ matrix.target }}
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Install Protoc v21.12
|
||||
run: choco install --no-progress protoc
|
||||
- name: Build
|
||||
@@ -208,11 +240,12 @@ jobs:
|
||||
$env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT
|
||||
cargo build --profile ci --features aws,remote --tests --locked --target ${{ matrix.target }}
|
||||
- name: Run tests
|
||||
# Can only run tests when target matches host
|
||||
if: ${{ matrix.target == 'x86_64-pc-windows-msvc' }}
|
||||
run: |
|
||||
$env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT
|
||||
cargo test --profile ci --features aws,remote --locked
|
||||
# `--target` has to match the build step above. Without it cargo uses
|
||||
# target/ci/ rather than target/<triple>/ci/ and rebuilds the entire
|
||||
# dependency graph a second time.
|
||||
cargo test --profile ci --features aws,remote --locked --target ${{ matrix.target }}
|
||||
|
||||
msrv:
|
||||
# Check the minimum supported Rust version
|
||||
@@ -238,6 +271,11 @@ jobs:
|
||||
with:
|
||||
toolchain: ${{ matrix.msrv }}
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# Restore everywhere, but only save from main. Per-PR saves are
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Downgrade dependencies
|
||||
# These packages have newer requirements for MSRV
|
||||
run: |
|
||||
@@ -258,16 +296,18 @@ jobs:
|
||||
cargo update -p aws-types --precise 1.3.9
|
||||
cargo update -p aws-sigv4 --precise 1.3.5
|
||||
cargo update -p aws-credential-types --precise 1.2.8
|
||||
cargo update -p aws-smithy-checksums --precise 0.63.9
|
||||
# aws-smithy-checksums must stay at or above 0.63.13: OpenDAL's S3
|
||||
# service needs crc-fast ~1.9, and older releases pin it to ~1.3.
|
||||
cargo update -p aws-smithy-checksums --precise 0.63.13
|
||||
cargo update -p aws-smithy-runtime --precise 1.9.3
|
||||
cargo update -p aws-smithy-http --precise 0.62.4
|
||||
cargo update -p aws-smithy-eventstream --precise 0.60.12
|
||||
cargo update -p aws-smithy-http --precise 0.62.6
|
||||
cargo update -p aws-smithy-eventstream --precise 0.60.14
|
||||
cargo update -p aws-smithy-http-client --precise 1.1.3
|
||||
cargo update -p aws-smithy-observability --precise 0.1.4
|
||||
cargo update -p aws-smithy-query --precise 0.60.8
|
||||
cargo update -p aws-smithy-runtime-api --precise 1.9.1
|
||||
cargo update -p aws-smithy-async --precise 1.2.6
|
||||
cargo update -p aws-smithy-types --precise 1.3.5
|
||||
cargo update -p aws-smithy-runtime-api --precise 1.9.3
|
||||
cargo update -p aws-smithy-async --precise 1.2.7
|
||||
cargo update -p aws-smithy-types --precise 1.3.6
|
||||
cargo update -p aws-smithy-xml --precise 0.60.11
|
||||
cargo update -p home --precise 0.5.9
|
||||
- name: cargo +${{ matrix.msrv }} check
|
||||
|
||||
@@ -92,6 +92,8 @@ Python bindings changes:
|
||||
* Should use `LOOP.run()` to call the corresponding `AsyncTable` method.
|
||||
6. Add concrete sync method to `RemoteTable` class in `python/python/lancedb/remote/table.py`.
|
||||
7. Add unit test in `python/tests/test_table.py`.
|
||||
8. If you added a new public class or module-level function (not just a method on an
|
||||
existing class), expose it in the API reference. See "Python API reference" below.
|
||||
|
||||
TypeScript bindings changes:
|
||||
|
||||
@@ -103,6 +105,33 @@ TypeScript bindings changes:
|
||||
5. Add test in `nodejs/__test__/table.test.ts`.
|
||||
6. Run `npm run docs` to generate TypeScript documentation.
|
||||
|
||||
## Python API reference
|
||||
|
||||
`docs/src/python/python.md` is the entire Python API reference. It is maintained by
|
||||
hand, and anything not listed there is not rendered at all, so new public classes and
|
||||
module-level functions have to be added explicitly. How depends on the module:
|
||||
|
||||
* `lancedb.index`, `lancedb.embeddings`, `lancedb.remote`, and `lancedb.rerankers` are
|
||||
rendered by a single directive each, driven by the module's `__all__`. Add the new
|
||||
name to `__all__` and it appears; forget, and it is silently omitted.
|
||||
* Everything else (`lancedb`, `lancedb.table`, `lancedb.query`, `lancedb.db`, ...) is
|
||||
listed symbol by symbol. Add a `::: lancedb.<module>.<Name>` line to the matching
|
||||
section, and remember that the page separates synchronous and asynchronous APIs.
|
||||
|
||||
Deliberately undocumented: concrete implementations reached through an abstract base
|
||||
(`LanceTable`, `LanceDBConnection`, `RemoteDBConnection`), query base classes already
|
||||
covered by `inherited_members`, and internal helpers.
|
||||
|
||||
Cross-references in docstrings use mkdocstrings syntax, `[text][lancedb.table.Table]`.
|
||||
Plain relative links such as `[Table](Table)` do not resolve. To check your work:
|
||||
|
||||
```shell
|
||||
pip install -r docs/requirements.txt
|
||||
cd docs && PYTHONPATH=. mkdocs build
|
||||
```
|
||||
|
||||
The docs site only builds on pushes to `main`, so this is not covered by PR CI.
|
||||
|
||||
## Review Guidelines
|
||||
|
||||
Please consider the following when reviewing code contributions.
|
||||
|
||||
Generated
+421
-354
File diff suppressed because it is too large
Load Diff
+16
-15
@@ -1,5 +1,6 @@
|
||||
[workspace]
|
||||
members = ["rust/lancedb", "nodejs", "python"]
|
||||
exclude = ["lance-artifact"]
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
@@ -13,20 +14,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=9.1.0-beta.8", default-features = false, "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.1.0-beta.8", default-features = false, "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.1.0-beta.8", default-features = false, "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.1.0-beta.8", "tag" = "v9.1.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a", default-features = false }
|
||||
lance-core = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-datagen = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-file = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-io = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a", default-features = false }
|
||||
lance-index = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-linalg = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-namespace = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-namespace-impls = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a", default-features = false }
|
||||
lance-table = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-testing = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-datafusion = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-encoding = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
lance-arrow = { version = "=11.0.0-beta.3", git = "https://github.com/lancedb/lancedb.git", rev = "31b68dd21b60fcbee4354a175c9d65ff6ae02b6a" }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
@@ -52,7 +53,7 @@ env_logger = "0.11"
|
||||
half = { "version" = "2.7.1", default-features = false, features = [
|
||||
"num-traits",
|
||||
] }
|
||||
futures = "0"
|
||||
futures = "0.3"
|
||||
log = "0.4"
|
||||
metrics = "0.24"
|
||||
metrics-util = "0.19"
|
||||
|
||||
+3
-3
@@ -2,9 +2,9 @@ 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)}
|
||||
HEAD_SHA=$(git rev-parse HEAD)
|
||||
|
||||
readonly TAG_PREFIX="v"
|
||||
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)
|
||||
@@ -12,7 +12,7 @@ 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
|
||||
if [[ "$RELEASE_TYPE" == 'stable' ]]; then
|
||||
BUMP_ARGS="--no-tag"
|
||||
fi
|
||||
|
||||
|
||||
@@ -51,6 +51,11 @@ plugins:
|
||||
paths: [../python/python]
|
||||
options:
|
||||
docstring_style: numpy
|
||||
docstring_options:
|
||||
# Attributes documented in a `Parameters` section, and pydantic
|
||||
# dataclasses whose `__init__` griffe cannot see statically, both
|
||||
# trip this check. It reports nothing actionable here.
|
||||
warn_unknown_params: false
|
||||
heading_level: 3
|
||||
show_signature_annotations: true
|
||||
show_root_heading: true
|
||||
|
||||
+11
-1
@@ -453,6 +453,16 @@ paths:
|
||||
The metric type to use for the index. l2, Cosine, Dot are supported.
|
||||
index_type:
|
||||
type: string
|
||||
custom_stop_words:
|
||||
type: [array, "null"]
|
||||
items:
|
||||
type: string
|
||||
description: |
|
||||
The custom stop-word list for an FTS index. A non-null
|
||||
array replaces the language's built-in stop-word list and is only
|
||||
applied when remove_stop_words is enabled. Null uses the built-in
|
||||
language list, while an empty array explicitly replaces it with no
|
||||
stop words.
|
||||
responses:
|
||||
"200":
|
||||
description: Index successfully created
|
||||
@@ -510,4 +520,4 @@ paths:
|
||||
"401":
|
||||
$ref: "#/components/responses/unauthorized"
|
||||
"404":
|
||||
$ref: "#/components/responses/not_found"
|
||||
$ref: "#/components/responses/not_found"
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<version>0.37.1-beta.0</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Contributing to LanceDB Typescript
|
||||
|
||||
This document outlines the process for contributing to LanceDB Typescript.
|
||||
For general contribution guidelines, see [CONTRIBUTING.md](../CONTRIBUTING.md).
|
||||
For general contribution guidelines, see [CONTRIBUTING.md](https://github.com/lancedb/lancedb/blob/main/CONTRIBUTING.md).
|
||||
|
||||
## Project layout
|
||||
|
||||
|
||||
@@ -25,6 +25,27 @@ the underlying connection has been closed.
|
||||
|
||||
## Methods
|
||||
|
||||
### cancelJob()
|
||||
|
||||
```ts
|
||||
abstract cancelJob(jobId): Promise<boolean>
|
||||
```
|
||||
|
||||
Request cancellation of a server-side job by id.
|
||||
|
||||
Resolves to true if the server accepted the cancellation, false if no
|
||||
such job exists. Cancelling an already-terminal job is a no-op success.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`boolean`>
|
||||
|
||||
***
|
||||
|
||||
### cloneTable()
|
||||
|
||||
```ts
|
||||
@@ -365,6 +386,26 @@ Drop an existing table.
|
||||
|
||||
***
|
||||
|
||||
### getJob()
|
||||
|
||||
```ts
|
||||
abstract getJob(jobId): Promise<null | JobDescription>
|
||||
```
|
||||
|
||||
Describe a single server-side job by id.
|
||||
|
||||
Resolves to `null` when the server has no such job.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`null` \| [`JobDescription`](../interfaces/JobDescription.md)>
|
||||
|
||||
***
|
||||
|
||||
### isOpen()
|
||||
|
||||
```ts
|
||||
@@ -379,6 +420,62 @@ Return true if the connection has not been closed
|
||||
|
||||
***
|
||||
|
||||
### job()
|
||||
|
||||
```ts
|
||||
abstract job(jobId): Job
|
||||
```
|
||||
|
||||
A [Job](Job.md) handle for a server-side job by id.
|
||||
|
||||
The handle is constructed without a server round trip; an unknown id
|
||||
surfaces when the handle is used. Dropping the handle has no effect on
|
||||
the job itself.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
[`Job`](Job.md)
|
||||
|
||||
***
|
||||
|
||||
### jobHistory()
|
||||
|
||||
```ts
|
||||
abstract jobHistory(jobId?): Promise<Table<any>>
|
||||
```
|
||||
|
||||
The lifecycle event history of a server-side job, as an Arrow table.
|
||||
|
||||
Lists history across all jobs when `jobId` is omitted.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId?**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Table`<`any`>>
|
||||
|
||||
***
|
||||
|
||||
### listJobs()
|
||||
|
||||
```ts
|
||||
abstract listJobs(): Promise<JobInfo[]>
|
||||
```
|
||||
|
||||
List server-side jobs across the database's tables.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`JobInfo`](../interfaces/JobInfo.md)[]>
|
||||
|
||||
***
|
||||
|
||||
### listNamespaces()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / Job
|
||||
|
||||
# Class: Job
|
||||
|
||||
A handle to an operation that may still be running.
|
||||
|
||||
## Constructors
|
||||
|
||||
### new Job()
|
||||
|
||||
```ts
|
||||
new Job(): Job
|
||||
```
|
||||
|
||||
#### Returns
|
||||
|
||||
[`Job`](Job.md)
|
||||
|
||||
## Accessors
|
||||
|
||||
### id
|
||||
|
||||
```ts
|
||||
get id(): null | string
|
||||
```
|
||||
|
||||
Identifies the operation on the server that is running it. Operations
|
||||
that run in this process have no server id. The value is opaque.
|
||||
|
||||
#### Returns
|
||||
|
||||
`null` \| `string`
|
||||
|
||||
## Methods
|
||||
|
||||
### cancel()
|
||||
|
||||
```ts
|
||||
cancel(): Promise<void>
|
||||
```
|
||||
|
||||
Request cancellation. Cancelling a finished operation is a no-op.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`void`>
|
||||
|
||||
***
|
||||
|
||||
### status()
|
||||
|
||||
```ts
|
||||
status(): Promise<string>
|
||||
```
|
||||
|
||||
The operation's current lifecycle state: "running", "finished",
|
||||
"failed", or "cancelled".
|
||||
|
||||
A point snapshot; unlike [Job.wait](Job.md#wait) it does not block or reject
|
||||
on a terminal failure state. States a newer server reports that this
|
||||
client version does not know pass through as-is.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`string`>
|
||||
|
||||
***
|
||||
|
||||
### wait()
|
||||
|
||||
```ts
|
||||
wait(): Promise<void>
|
||||
```
|
||||
|
||||
Wait until the operation reaches a terminal state.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`void`>
|
||||
@@ -76,24 +76,23 @@ the query optimizer chooses a suboptimal path.
|
||||
|
||||
***
|
||||
|
||||
### useLsmWrite()
|
||||
### useLsm()
|
||||
|
||||
```ts
|
||||
useLsmWrite(useLsmWrite): MergeInsertBuilder
|
||||
useLsm(enable): MergeInsertBuilder
|
||||
```
|
||||
|
||||
Controls whether the merge uses the MemWAL LSM write path.
|
||||
Control MemWAL routing for this merge.
|
||||
|
||||
By default (unset), a `mergeInsert` on a table with an LSM write spec is
|
||||
routed through Lance's MemWAL shard writer, and a table without one uses
|
||||
the standard path. Pass `false` to force the standard path even when a
|
||||
spec is set. Pass `true` to require a spec — `mergeInsert` rejects if none
|
||||
is installed.
|
||||
routed through Lance's MemWAL shard writer, and a table without one uses the
|
||||
standard path.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **useLsmWrite**: `boolean`
|
||||
Whether to use the LSM write path.
|
||||
* **enable**: `boolean`
|
||||
`true` forces MemWAL routing and errors if the table has no
|
||||
LSM write spec. `false` forces the standard write path even when a spec is set.
|
||||
|
||||
#### Returns
|
||||
|
||||
|
||||
@@ -497,6 +497,42 @@ ArrowTable.
|
||||
|
||||
***
|
||||
|
||||
### useLsm()
|
||||
|
||||
```ts
|
||||
useLsm(enable): this
|
||||
```
|
||||
|
||||
Control MemWAL read routing for this query.
|
||||
|
||||
By default (unset), when the table carries a MemWAL write spec (see
|
||||
[Table#setLsmWriteSpec](Table.md#setlsmwritespec)), reads are routed through the LSM scanner so
|
||||
they also return data written via the `mergeInsert` LSM path that has not yet
|
||||
been compacted into the base table (the active/frozen in-memory memtables and
|
||||
the flushed generations), deduplicated by primary key; a table without a spec
|
||||
reads the base table.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **enable**: `boolean`
|
||||
`true` forces the LSM scanner and errors if the table has no
|
||||
MemWAL write spec. `false` bypasses the MemWAL and reads the base table only,
|
||||
even when a spec is present.
|
||||
Note: the LSM scanner does not support every query shape (e.g. reranking,
|
||||
hybrid search, `orderBy`). On a MemWAL table those shapes error unless
|
||||
`useLsm(false)` is set, because a base-only read would silently exclude
|
||||
un-compacted MemWAL data.
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.useLsm`
|
||||
|
||||
***
|
||||
|
||||
### where()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -295,6 +295,29 @@ await table.createIndex("my_float_col");
|
||||
|
||||
***
|
||||
|
||||
### createIndexAsync()
|
||||
|
||||
```ts
|
||||
abstract createIndexAsync(column, options?): Promise<Job>
|
||||
```
|
||||
|
||||
Create an index, returning a handle to the indexing job.
|
||||
|
||||
The job may already be complete when returned; callers must not assume
|
||||
the index exists until [Job.wait](Job.md#wait) resolves.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **column**: `string`
|
||||
|
||||
* **options?**: `Partial`<[`IndexOptions`](../interfaces/IndexOptions.md)>
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`Job`](Job.md)>
|
||||
|
||||
***
|
||||
|
||||
### currentBranch()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -273,6 +273,29 @@ ArrowTable.
|
||||
|
||||
***
|
||||
|
||||
### useLsm()
|
||||
|
||||
```ts
|
||||
useLsm(enable): this
|
||||
```
|
||||
|
||||
Control MemWAL read routing for this take query.
|
||||
|
||||
`false` bypasses the MemWAL and reads the base table only — the escape hatch,
|
||||
since take-by-row-id/offset is not supported on the LSM scanner and, on a
|
||||
MemWAL table, auto-routes to it and errors otherwise.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **enable**: `boolean`
|
||||
`false` reads the base table only.
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
***
|
||||
|
||||
### withRowId()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -746,6 +746,42 @@ ArrowTable.
|
||||
|
||||
***
|
||||
|
||||
### useLsm()
|
||||
|
||||
```ts
|
||||
useLsm(enable): this
|
||||
```
|
||||
|
||||
Control MemWAL read routing for this query.
|
||||
|
||||
By default (unset), when the table carries a MemWAL write spec (see
|
||||
[Table#setLsmWriteSpec](Table.md#setlsmwritespec)), reads are routed through the LSM scanner so
|
||||
they also return data written via the `mergeInsert` LSM path that has not yet
|
||||
been compacted into the base table (the active/frozen in-memory memtables and
|
||||
the flushed generations), deduplicated by primary key; a table without a spec
|
||||
reads the base table.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **enable**: `boolean`
|
||||
`true` forces the LSM scanner and errors if the table has no
|
||||
MemWAL write spec. `false` bypasses the MemWAL and reads the base table only,
|
||||
even when a spec is present.
|
||||
Note: the LSM scanner does not support every query shape (e.g. reranking,
|
||||
hybrid search, `orderBy`). On a MemWAL table those shapes error unless
|
||||
`useLsm(false)` is set, because a base-only read would silently exclude
|
||||
un-compacted MemWAL data.
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.useLsm`
|
||||
|
||||
***
|
||||
|
||||
### where()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -25,6 +25,7 @@
|
||||
- [Connection](classes/Connection.md)
|
||||
- [HeaderProvider](classes/HeaderProvider.md)
|
||||
- [Index](classes/Index.md)
|
||||
- [Job](classes/Job.md)
|
||||
- [MakeArrowTableOptions](classes/MakeArrowTableOptions.md)
|
||||
- [MatchQuery](classes/MatchQuery.md)
|
||||
- [MergeInsertBuilder](classes/MergeInsertBuilder.md)
|
||||
@@ -88,6 +89,9 @@
|
||||
- [IvfFlatOptions](interfaces/IvfFlatOptions.md)
|
||||
- [IvfPqOptions](interfaces/IvfPqOptions.md)
|
||||
- [IvfRqOptions](interfaces/IvfRqOptions.md)
|
||||
- [JobDescription](interfaces/JobDescription.md)
|
||||
- [JobFailureInfo](interfaces/JobFailureInfo.md)
|
||||
- [JobInfo](interfaces/JobInfo.md)
|
||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
||||
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
|
||||
|
||||
@@ -43,6 +43,34 @@ The following tokenizers are available:
|
||||
|
||||
***
|
||||
|
||||
### blockSize?
|
||||
|
||||
```ts
|
||||
optional blockSize: 128 | 256;
|
||||
```
|
||||
|
||||
Number of documents per compressed posting block.
|
||||
|
||||
The default is 128. Supported values are 128 and 256. A value of 256 uses
|
||||
the experimental FTS V3 format and may introduce breaking changes.
|
||||
|
||||
***
|
||||
|
||||
### customStopWords?
|
||||
|
||||
```ts
|
||||
optional customStopWords: string[];
|
||||
```
|
||||
|
||||
Custom stop words that replace the built-in list for `language`.
|
||||
|
||||
This option only affects tokenization when `removeStopWords` is true.
|
||||
|
||||
`undefined` keeps the built-in language list. An empty array explicitly
|
||||
replaces it with no stop words.
|
||||
|
||||
***
|
||||
|
||||
### language?
|
||||
|
||||
```ts
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / JobDescription
|
||||
|
||||
# Interface: JobDescription
|
||||
|
||||
A described job from `Connection.getJob`.
|
||||
|
||||
## Properties
|
||||
|
||||
### creationMs
|
||||
|
||||
```ts
|
||||
creationMs: number;
|
||||
```
|
||||
|
||||
When the job was created, in milliseconds since the epoch.
|
||||
|
||||
***
|
||||
|
||||
### failure?
|
||||
|
||||
```ts
|
||||
optional failure: JobFailureInfo;
|
||||
```
|
||||
|
||||
Why the job failed, when the job is failed and the server reports a
|
||||
reason.
|
||||
|
||||
***
|
||||
|
||||
### jobId
|
||||
|
||||
```ts
|
||||
jobId: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### jobType
|
||||
|
||||
```ts
|
||||
jobType: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### specJson?
|
||||
|
||||
```ts
|
||||
optional specJson: string;
|
||||
```
|
||||
|
||||
The job-type-specific specification as a JSON string, when present.
|
||||
|
||||
***
|
||||
|
||||
### state
|
||||
|
||||
```ts
|
||||
state: string;
|
||||
```
|
||||
|
||||
Lifecycle state: "running", "finished", "failed", or "cancelled".
|
||||
@@ -0,0 +1,33 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / JobFailureInfo
|
||||
|
||||
# Interface: JobFailureInfo
|
||||
|
||||
The server's account of why a job failed.
|
||||
|
||||
## Properties
|
||||
|
||||
### message?
|
||||
|
||||
```ts
|
||||
optional message: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### phase?
|
||||
|
||||
```ts
|
||||
optional phase: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### retryable?
|
||||
|
||||
```ts
|
||||
optional retryable: boolean;
|
||||
```
|
||||
@@ -0,0 +1,58 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / JobInfo
|
||||
|
||||
# Interface: JobInfo
|
||||
|
||||
A row from `Connection.listJobs`: one server-side job.
|
||||
|
||||
## Properties
|
||||
|
||||
### createdAtMillis
|
||||
|
||||
```ts
|
||||
createdAtMillis: number;
|
||||
```
|
||||
|
||||
When the job was created, in milliseconds since the epoch.
|
||||
|
||||
***
|
||||
|
||||
### jobId
|
||||
|
||||
```ts
|
||||
jobId: string;
|
||||
```
|
||||
|
||||
The job id -- what `Connection.getJob` and `Connection.cancelJob`
|
||||
accept.
|
||||
|
||||
***
|
||||
|
||||
### jobType
|
||||
|
||||
```ts
|
||||
jobType: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### state
|
||||
|
||||
```ts
|
||||
state: string;
|
||||
```
|
||||
|
||||
Lifecycle state: "running", "finished", "failed", or "cancelled".
|
||||
|
||||
***
|
||||
|
||||
### table
|
||||
|
||||
```ts
|
||||
table: string;
|
||||
```
|
||||
|
||||
The table the job runs against, without URI or namespace.
|
||||
@@ -30,6 +30,21 @@ The tokenizer to use. The default is "simple".
|
||||
|
||||
***
|
||||
|
||||
### customStopWords?
|
||||
|
||||
```ts
|
||||
optional customStopWords: string[];
|
||||
```
|
||||
|
||||
Custom stop words that replace the built-in list for `language`.
|
||||
|
||||
This option only affects tokenization when `removeStopWords` is true.
|
||||
|
||||
`undefined` keeps the built-in language list. An empty array explicitly
|
||||
replaces it with no stop words.
|
||||
|
||||
***
|
||||
|
||||
### language?
|
||||
|
||||
```ts
|
||||
|
||||
+141
-52
@@ -26,6 +26,18 @@ is also an [asynchronous API client](#connections-asynchronous).
|
||||
|
||||
::: lancedb.db.DBConnection
|
||||
|
||||
::: lancedb.Session
|
||||
|
||||
## Namespaces (Synchronous)
|
||||
|
||||
A namespace-backed connection resolves tables through a
|
||||
[Lance namespace](https://lance-format.github.io/lance-namespace/) service instead of
|
||||
listing a storage directory.
|
||||
|
||||
::: lancedb.connect_namespace
|
||||
|
||||
::: lancedb.namespace.LanceNamespaceDBConnection
|
||||
|
||||
## Tables (Synchronous)
|
||||
|
||||
::: lancedb.table.Table
|
||||
@@ -34,8 +46,12 @@ is also an [asynchronous API client](#connections-asynchronous).
|
||||
|
||||
::: lancedb.table.FragmentSummaryStats
|
||||
|
||||
::: lancedb.table.TableStatistics
|
||||
|
||||
::: lancedb.table.Tags
|
||||
|
||||
::: lancedb.table.Branches
|
||||
|
||||
## Expressions
|
||||
|
||||
Type-safe expression builder for filters and projections. Use these instead
|
||||
@@ -62,29 +78,46 @@ of raw SQL strings with [where][lancedb.query.LanceQueryBuilder.where] and
|
||||
|
||||
::: lancedb.query.LanceHybridQueryBuilder
|
||||
|
||||
::: lancedb.query.LanceEmptyQueryBuilder
|
||||
|
||||
::: lancedb.query.LanceTakeQueryBuilder
|
||||
|
||||
## Full text queries
|
||||
|
||||
Structured full text queries can be passed to
|
||||
[Table.search][lancedb.table.Table.search] or
|
||||
[AsyncTable.search][lancedb.table.AsyncTable.search] in place of a query string,
|
||||
and combined with [BooleanQuery][lancedb.query.BooleanQuery].
|
||||
|
||||
::: lancedb.query.FullTextQuery
|
||||
|
||||
::: lancedb.query.MatchQuery
|
||||
|
||||
::: lancedb.query.PhraseQuery
|
||||
|
||||
::: lancedb.query.BoostQuery
|
||||
|
||||
::: lancedb.query.MultiMatchQuery
|
||||
|
||||
::: lancedb.query.BooleanQuery
|
||||
|
||||
::: lancedb.query.FullTextOperator
|
||||
|
||||
::: lancedb.query.Occur
|
||||
|
||||
## Embeddings
|
||||
|
||||
::: lancedb.embeddings.registry.EmbeddingFunctionRegistry
|
||||
|
||||
::: lancedb.embeddings.base.EmbeddingFunctionConfig
|
||||
|
||||
::: lancedb.embeddings.base.EmbeddingFunction
|
||||
|
||||
::: lancedb.embeddings.base.TextEmbeddingFunction
|
||||
|
||||
::: lancedb.embeddings.sentence_transformers.SentenceTransformerEmbeddings
|
||||
|
||||
::: lancedb.embeddings.openai.OpenAIEmbeddings
|
||||
|
||||
::: lancedb.embeddings.open_clip.OpenClipEmbeddings
|
||||
::: lancedb.embeddings
|
||||
options:
|
||||
show_root_heading: false
|
||||
show_root_toc_entry: false
|
||||
|
||||
## Remote configuration
|
||||
|
||||
::: lancedb.remote.ClientConfig
|
||||
|
||||
::: lancedb.remote.TimeoutConfig
|
||||
|
||||
::: lancedb.remote.RetryConfig
|
||||
::: lancedb.remote
|
||||
options:
|
||||
show_root_heading: false
|
||||
show_root_toc_entry: false
|
||||
|
||||
## Context
|
||||
|
||||
@@ -94,11 +127,50 @@ of raw SQL strings with [where][lancedb.query.LanceQueryBuilder.where] and
|
||||
|
||||
## Full text search
|
||||
|
||||
Use [lancedb.table.Table.create_fts_index][] for the synchronous API or
|
||||
[lancedb.table.AsyncTable.create_index][] with [lancedb.index.FTS][] for the
|
||||
asynchronous API.
|
||||
Pass `custom_stop_words` to [lancedb.index.FTS][]:
|
||||
|
||||
::: lancedb.index.FTS
|
||||
```python
|
||||
from lancedb.index import FTS
|
||||
|
||||
table.create_index(
|
||||
"text",
|
||||
config=FTS(remove_stop_words=True, custom_stop_words=["acme", "internal"]),
|
||||
)
|
||||
```
|
||||
|
||||
The list replaces the built-in stop words and is used only when
|
||||
`remove_stop_words=True`:
|
||||
|
||||
- `custom_stop_words=None` uses the built-in list for `language`.
|
||||
- `custom_stop_words=[]` removes no words.
|
||||
- Values are passed through without trimming, lowercasing, or other rewriting.
|
||||
|
||||
The same option is available on `lancedb.tokenize(...)` and the deprecated
|
||||
[lancedb.table.Table.create_fts_index][] compatibility helper:
|
||||
|
||||
```python
|
||||
import lancedb
|
||||
|
||||
tokens = list(lancedb.tokenize("acme makes searchable data",
|
||||
custom_stop_words=["acme"]))
|
||||
```
|
||||
|
||||
::: lancedb.tokenize
|
||||
|
||||
::: lancedb.FtsToken
|
||||
|
||||
## Blobs
|
||||
|
||||
Blob columns store large binary values out of line so they can be read lazily
|
||||
instead of being materialized with the rest of the row.
|
||||
|
||||
::: lancedb.blob
|
||||
|
||||
::: lancedb.BlobType
|
||||
|
||||
::: lancedb._blob.BlobFile
|
||||
options:
|
||||
show_root_full_path: false
|
||||
|
||||
## Utilities
|
||||
|
||||
@@ -106,6 +178,14 @@ asynchronous API.
|
||||
|
||||
::: lancedb.merge.LanceMergeInsertBuilder
|
||||
|
||||
::: lancedb.otel.instrument_lancedb_metrics
|
||||
|
||||
## Exceptions
|
||||
|
||||
::: lancedb.exceptions.MissingValueError
|
||||
|
||||
::: lancedb.exceptions.MissingColumnError
|
||||
|
||||
## Integrations
|
||||
|
||||
## Pydantic
|
||||
@@ -114,19 +194,30 @@ asynchronous API.
|
||||
|
||||
::: lancedb.pydantic.vector
|
||||
|
||||
::: lancedb.pydantic.Vector
|
||||
|
||||
::: lancedb.pydantic.MultiVector
|
||||
|
||||
::: lancedb.pydantic.LanceModel
|
||||
|
||||
## PyTorch
|
||||
|
||||
::: lancedb.streaming.StreamingDataset
|
||||
|
||||
::: lancedb.permutation.permutation_builder
|
||||
|
||||
::: lancedb.permutation.PermutationBuilder
|
||||
|
||||
::: lancedb.permutation.Permutation
|
||||
|
||||
::: lancedb.permutation.Transforms
|
||||
|
||||
## Reranking
|
||||
|
||||
::: lancedb.rerankers.linear_combination.LinearCombinationReranker
|
||||
|
||||
::: lancedb.rerankers.cohere.CohereReranker
|
||||
|
||||
::: lancedb.rerankers.colbert.ColbertReranker
|
||||
|
||||
::: lancedb.rerankers.cross_encoder.CrossEncoderReranker
|
||||
|
||||
::: lancedb.rerankers.openai.OpenaiReranker
|
||||
::: lancedb.rerankers
|
||||
options:
|
||||
show_root_heading: false
|
||||
show_root_toc_entry: false
|
||||
|
||||
## Connections (Asynchronous)
|
||||
|
||||
@@ -137,6 +228,12 @@ can be used to create, list, or open tables.
|
||||
|
||||
::: lancedb.db.AsyncConnection
|
||||
|
||||
## Namespaces (Asynchronous)
|
||||
|
||||
::: lancedb.connect_namespace_async
|
||||
|
||||
::: lancedb.namespace.AsyncLanceNamespaceDBConnection
|
||||
|
||||
## Tables (Asynchronous)
|
||||
|
||||
Table hold your actual data as a collection of records / rows.
|
||||
@@ -145,32 +242,20 @@ Table hold your actual data as a collection of records / rows.
|
||||
|
||||
::: lancedb.table.AsyncTags
|
||||
|
||||
::: lancedb.table.AsyncBranches
|
||||
|
||||
## Indices (Asynchronous)
|
||||
|
||||
Indices can be created on a table to speed up queries. This section
|
||||
lists the indices that LanceDb supports.
|
||||
|
||||
::: lancedb.index.BTree
|
||||
|
||||
::: lancedb.index.Bitmap
|
||||
|
||||
::: lancedb.index.LabelList
|
||||
|
||||
::: lancedb.index.FTS
|
||||
|
||||
::: lancedb.index.IvfPq
|
||||
|
||||
::: lancedb.index.HnswPq
|
||||
|
||||
::: lancedb.index.HnswSq
|
||||
|
||||
::: lancedb.index.IvfFlat
|
||||
|
||||
::: lancedb.index.IvfSq
|
||||
|
||||
::: lancedb.index.IvfRq
|
||||
|
||||
::: lancedb.index.HnswFlat
|
||||
::: lancedb.index
|
||||
options:
|
||||
show_root_heading: false
|
||||
show_root_toc_entry: false
|
||||
# `lang_mapping` is defined in the module rather than imported, so it is
|
||||
# picked up despite not being in `__all__`. It is an internal lookup table.
|
||||
filters: ["!^_", "!^lang_mapping$"]
|
||||
|
||||
::: lancedb.table.IndexStatistics
|
||||
|
||||
@@ -198,3 +283,7 @@ rows nearest to a query vector and can be created with the
|
||||
::: lancedb.query.AsyncHybridQuery
|
||||
options:
|
||||
inherited_members: true
|
||||
|
||||
::: lancedb.query.AsyncTakeQuery
|
||||
options:
|
||||
inherited_members: true
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<version>0.37.1-beta.0</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<version>0.37.1-beta.0</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>${project.artifactId}</name>
|
||||
<description>LanceDB Java SDK Parent POM</description>
|
||||
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>9.1.0-beta.8</lance-core.version>
|
||||
<lance-core.version>11.0.0-beta.3</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
[workspace]
|
||||
members = [
|
||||
"third-party/opendal",
|
||||
"third-party/opendal-service-s3",
|
||||
"rust/examples",
|
||||
"rust/lance",
|
||||
"rust/lance-arrow",
|
||||
"rust/lance-core",
|
||||
"rust/lance-datagen",
|
||||
"rust/lance-encoding",
|
||||
"rust/lance-file",
|
||||
"rust/lance-geo",
|
||||
"rust/lance-index",
|
||||
"rust/lance-index-core",
|
||||
"rust/lance-io",
|
||||
"rust/lance-linalg",
|
||||
"rust/lance-namespace",
|
||||
"rust/lance-namespace-impls",
|
||||
"rust/lance-namespace-datafusion",
|
||||
"rust/lance-select",
|
||||
"rust/lance-tokenizer",
|
||||
"rust/lance-table",
|
||||
"rust/lance-derive",
|
||||
"rust/lance-test-macros",
|
||||
"rust/lance-testing",
|
||||
"rust/lance-tools",
|
||||
"rust/compression/fsst",
|
||||
"rust/compression/bitpacking",
|
||||
"rust/arrow-scalar",
|
||||
"rust/arrow-stats",
|
||||
]
|
||||
exclude = ["python", "java/lance-jni"]
|
||||
# Python package needs to be built by maturin.
|
||||
resolver = "3"
|
||||
|
||||
|
||||
[workspace.package]
|
||||
version = "11.0.0-beta.3"
|
||||
edition = "2024"
|
||||
authors = ["Lance Devs <dev@lance.org>"]
|
||||
license = "Apache-2.0"
|
||||
repository = "https://github.com/lance-format/lance"
|
||||
readme = "README.md"
|
||||
description = "A columnar data format that is 100x faster than Parquet for random access."
|
||||
keywords = [
|
||||
"data-format",
|
||||
"data-science",
|
||||
"machine-learning",
|
||||
"apache-arrow",
|
||||
"data-analytics",
|
||||
]
|
||||
categories = [
|
||||
"database-implementations",
|
||||
"data-structures",
|
||||
"development-tools",
|
||||
"science",
|
||||
]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
arc-swap = "1.7"
|
||||
libc = "0.2.176"
|
||||
lance = { version = "=11.0.0-beta.3", path = "./rust/lance", default-features = false }
|
||||
lance-arrow = { version = "=11.0.0-beta.3", path = "./rust/lance-arrow" }
|
||||
lance-core = { version = "=11.0.0-beta.3", path = "./rust/lance-core" }
|
||||
lance-datafusion = { version = "=11.0.0-beta.3", path = "./rust/lance-datafusion" }
|
||||
lance-datagen = { version = "=11.0.0-beta.3", path = "./rust/lance-datagen" }
|
||||
lance-derive = { version = "=11.0.0-beta.3", path = "./rust/lance-derive" }
|
||||
lance-encoding = { version = "=11.0.0-beta.3", path = "./rust/lance-encoding" }
|
||||
lance-file = { version = "=11.0.0-beta.3", path = "./rust/lance-file" }
|
||||
lance-geo = { version = "=11.0.0-beta.3", path = "./rust/lance-geo" }
|
||||
lance-index = { version = "=11.0.0-beta.3", path = "./rust/lance-index" }
|
||||
lance-index-core = { version = "=11.0.0-beta.3", path = "./rust/lance-index-core" }
|
||||
lance-io = { version = "=11.0.0-beta.3", path = "./rust/lance-io", default-features = false }
|
||||
lance-linalg = { version = "=11.0.0-beta.3", path = "./rust/lance-linalg" }
|
||||
lance-namespace = { version = "=11.0.0-beta.3", path = "./rust/lance-namespace" }
|
||||
lance-namespace-impls = { version = "=11.0.0-beta.3", path = "./rust/lance-namespace-impls" }
|
||||
lance-namespace-datafusion = { version = "=7.0.0-beta.9", path = "./rust/lance-namespace-datafusion" }
|
||||
lance-namespace-reqwest-client = "0.8.6"
|
||||
lance-select = { version = "=11.0.0-beta.3", path = "./rust/lance-select" }
|
||||
lance-tokenizer = { version = "=11.0.0-beta.3", path = "./rust/lance-tokenizer" }
|
||||
lance-table = { version = "=11.0.0-beta.3", path = "./rust/lance-table" }
|
||||
lance-test-macros = { version = "=11.0.0-beta.3", path = "./rust/lance-test-macros" }
|
||||
lance-testing = { version = "=11.0.0-beta.3", path = "./rust/lance-testing" }
|
||||
approx = "0.5.1"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false, features = ["prettyprint"] }
|
||||
lance-arrow-scalar = { version = "=58.0.0", path = "./rust/arrow-scalar" }
|
||||
lance-arrow-stats = { version = "=58.0.0", path = "./rust/arrow-stats" }
|
||||
arrow-arith = "58.0.0"
|
||||
arrow-array = "58.0.0"
|
||||
arrow-buffer = "58.0.0"
|
||||
arrow-cast = "58.0.0"
|
||||
arrow-data = "58.0.0"
|
||||
arrow-ipc = { version = "58.0.0", features = ["zstd"] }
|
||||
arrow-ord = "58.0.0"
|
||||
arrow-row = "58.0.0"
|
||||
arrow-schema = "58.0.0"
|
||||
arrow-select = "58.0.0"
|
||||
async-recursion = "1.0"
|
||||
async-trait = "0.1"
|
||||
axum = "0.7"
|
||||
aws-config = "1.2.0"
|
||||
aws-credential-types = "1.2.0"
|
||||
aws-sdk-dynamodb = { version = "1.38.0", default-features = false }
|
||||
aws-sdk-s3 = { version = "1.38.0", default-features = false }
|
||||
half = { "version" = "2.1", default-features = false, features = [
|
||||
"num-traits",
|
||||
"std",
|
||||
"bytemuck",
|
||||
] }
|
||||
lance-bitpacking = { version = "=11.0.0-beta.3", path = "./rust/compression/bitpacking" }
|
||||
bitpacking = "0.9"
|
||||
bitvec = "1"
|
||||
blake3 = "1.8.5"
|
||||
bytemuck = { version = "1", default-features = false, features = [
|
||||
"extern_crate_alloc",
|
||||
] }
|
||||
bytes = "1.11.1"
|
||||
byteorder = "1.5"
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
chrono = { version = "0.4.41", default-features = false, features = [
|
||||
"std",
|
||||
"now",
|
||||
"serde",
|
||||
] }
|
||||
criterion = { version = "0.8.2", features = [
|
||||
"async",
|
||||
"async_tokio",
|
||||
"html_reports",
|
||||
] }
|
||||
crossbeam-queue = "0.3"
|
||||
crossbeam-skiplist = "0.1"
|
||||
datafusion = { version = "54.0.0", default-features = false, features = [
|
||||
"crypto_expressions",
|
||||
"datetime_expressions",
|
||||
"encoding_expressions",
|
||||
"nested_expressions",
|
||||
"regex_expressions",
|
||||
"sql",
|
||||
"string_expressions",
|
||||
"unicode_expressions",
|
||||
] }
|
||||
datafusion-common = "54.0.0"
|
||||
datafusion-functions = { version = "54.0.0", default-features = false, features = ["regex_expressions"] }
|
||||
datafusion-sql = "54.0.0"
|
||||
datafusion-expr = "54.0.0"
|
||||
datafusion-ffi = "54.0.0"
|
||||
datafusion-physical-expr = "54.0.0"
|
||||
datafusion-physical-plan = "54.0.0"
|
||||
datafusion-substrait = { version = "54.0.0", default-features = false }
|
||||
dirs = "6.0.0"
|
||||
either = "1.0"
|
||||
fst = { version = "0.4.7", features = ["levenshtein"] }
|
||||
fsst = { version = "=11.0.0-beta.3", path = "./rust/compression/fsst" }
|
||||
futures = "0.3"
|
||||
geoarrow-array = "0.8"
|
||||
geoarrow-schema = "0.8"
|
||||
geodatafusion = "0.5.0"
|
||||
geo-traits = "0.3.0"
|
||||
geo-types = "0.7.16"
|
||||
http = "1.1.0"
|
||||
humantime = "2.2.0"
|
||||
hyperloglogplus = { version = "0.4.1", features = ["const-loop"] }
|
||||
icu_segmenter = { version = "2.2", default-features = false, features = ["compiled_data"] }
|
||||
io-uring = "0.7"
|
||||
itertools = "0.14"
|
||||
jieba-rs = { version = "0.10.0", default-features = false }
|
||||
jsonb = { version = "0.5.3", default-features = false, features = ["databend"] }
|
||||
libm = "0.2.15"
|
||||
log = "0.4"
|
||||
metrics = { version = "0.24" }
|
||||
metrics-util = { version = "0.19" }
|
||||
mockall = { version = "0.14.0" }
|
||||
mock_instant = { version = "0.6.0" }
|
||||
moka = { version = "0.12", features = ["future", "sync"] }
|
||||
ndarray = { version = "0.16.1", features = ["matrixmultiply-threading"] }
|
||||
num-traits = "0.2"
|
||||
object_store = { version = "0.13.2" }
|
||||
opendal = { version = "0.58.1", path = "./third-party/opendal" }
|
||||
object_store_opendal = { version = "0.58" }
|
||||
reqsign-aws-v4 = { version = "3.1.0" }
|
||||
reqsign-core = { version = "3.2.1" }
|
||||
reqsign-file-read-tokio = { version = "3.0.4" }
|
||||
pin-project = "1.0"
|
||||
path_abs = "0.5"
|
||||
pprof = { version = "0.15.0", features = ["flamegraph"] }
|
||||
proptest = "1.3.1"
|
||||
prost = "0.14.1"
|
||||
prost-build = "0.14.1"
|
||||
prost-types = "0.14.1"
|
||||
rand = { version = "0.9.1", features = ["small_rng"] }
|
||||
rand_distr = { version = "0.5.1" }
|
||||
rand_xoshiro = "0.7.0"
|
||||
rangemap = { version = "1.0" }
|
||||
rayon = "1.10"
|
||||
regex-syntax = "0.8.10"
|
||||
roaring = "0.11.4"
|
||||
rstest = "0.26.1"
|
||||
serde = { version = "^1" }
|
||||
serde_json = { version = "1" }
|
||||
semver = "1.0"
|
||||
serial_test = "3"
|
||||
snafu = "0.9"
|
||||
lindera = { version = "3.0.7" }
|
||||
tempfile = "3"
|
||||
test-log = { version = "0.2.15" }
|
||||
tokio = { version = "1.23", features = [
|
||||
"rt-multi-thread",
|
||||
"macros",
|
||||
"fs",
|
||||
"sync",
|
||||
] }
|
||||
tokio-stream = "0.1.14"
|
||||
tokio-util = { version = "0.7.16" }
|
||||
tower = "0.5"
|
||||
tower-http = "0.5"
|
||||
tracing = "0.1"
|
||||
tracing-mock = { version = "=0.1.0-beta.3" }
|
||||
twox-hash = "2.0"
|
||||
url = "2.5.7"
|
||||
uuid = { version = "1.2", features = ["v4", "serde"] }
|
||||
wiremock = "0.6"
|
||||
pretty_assertions = "1.4.0"
|
||||
|
||||
[profile.bench]
|
||||
opt-level = 3
|
||||
debug = true
|
||||
strip = false
|
||||
|
||||
[profile.ci]
|
||||
debug = "line-tables-only"
|
||||
inherits = "dev"
|
||||
incremental = false
|
||||
|
||||
# This rule applies to every package except workspace members (dependencies
|
||||
# such as `arrow` and `tokio`). It disables debug info and related features on
|
||||
# dependencies so their binaries stay smaller, improving cache reuse.
|
||||
[profile.ci.package."*"]
|
||||
debug = false
|
||||
debug-assertions = false
|
||||
strip = "debuginfo"
|
||||
incremental = false
|
||||
|
||||
[workspace.lints.rust]
|
||||
unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage,coverage_nightly)'] }
|
||||
unsafe_op_in_unsafe_fn = "allow"
|
||||
|
||||
[workspace.lints.clippy]
|
||||
all = { level = "deny", priority = -1 }
|
||||
style = { level = "deny", priority = -1 }
|
||||
cargo = { level = "deny", priority = -1 }
|
||||
fallible_impl_from = "deny"
|
||||
manual_let_else = "deny"
|
||||
redundant_pub_crate = "deny"
|
||||
string_add_assign = "deny"
|
||||
string_add = "deny"
|
||||
string_lit_as_bytes = "deny"
|
||||
use_self = "deny"
|
||||
dbg_macro = "deny"
|
||||
trait_duplication_in_bounds = "deny"
|
||||
redundant_clone = "deny"
|
||||
# We should always use log instead of println
|
||||
print_stdout = "deny"
|
||||
print_stderr = "deny"
|
||||
# not too much we can do to avoid multiple crate versions
|
||||
multiple-crate-versions = "allow"
|
||||
# We use Vec<Range<u64>> in a lot of places and it is very common to use a single range in the vec.
|
||||
single_range_in_vec_init = "allow"
|
||||
large_futures = "deny"
|
||||
disallowed_macros = "deny"
|
||||
@@ -0,0 +1,22 @@
|
||||
# LanceDB patch provenance
|
||||
|
||||
This artifact vendors the Lance 11.0.0-beta.3 Rust workspace from Lance commit
|
||||
`f7d475539cefbd140cc46a828f3d843e68cd10f1`. The complete workspace keeps all mutually coupled
|
||||
Lance crates on one Cargo source identity when `lancedb` consumes the pinned artifact commit.
|
||||
|
||||
The local patch makes AWS credential-family merging atomic before backend selection and teaches
|
||||
the built-in OpenDAL S3 signer to resolve credential-only storage options at request time. This
|
||||
keeps long-lived multipart uploads refreshable without rebuilding a store or changing its outer
|
||||
metadata. Keeping the change inside `AwsStoreProvider` leaves arbitrary registry providers and
|
||||
their complete `ObjectStore` results untouched. Remove this patch when the same behavior is
|
||||
available in the pinned Lance release.
|
||||
|
||||
The LanceDB workspace pins every coupled Lance crate to the immutable repository commit containing
|
||||
this artifact. That durable source survives transitive Git consumption instead of relying on a
|
||||
root `[patch]`, which Cargo ignores when LanceDB itself is used as a dependency.
|
||||
|
||||
The artifact also contains Apache OpenDAL 0.58.1's `opendal` and `opendal-service-s3` crates. The
|
||||
only OpenDAL change adds a direct custom-provider hook alongside its existing credential-chain
|
||||
hook. Lance uses the direct hook so a selected dynamic authority can propagate refresh and
|
||||
validation errors; static or ambient credentials are considered only when that authority returns
|
||||
`Ok(None)`. Remove these copies when upstream OpenDAL exposes an equivalent hook.
|
||||
@@ -0,0 +1,255 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
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.
|
||||
|
||||
--------------------------------------------------------------------------------
|
||||
|
||||
This project includes code from Ritchie Vink's Polars project, which is licensed
|
||||
under the MIT license:
|
||||
|
||||
Copyright (c) 2020 Ritchie Vink
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
https://github.com/pola-rs/polars/blob/main/LICENSE
|
||||
|
||||
--------------------------------------------------------------------------------
|
||||
|
||||
This project includes code adapted from the quickwit-oss/bitpacking crate, which
|
||||
is licensed under the MIT license:
|
||||
|
||||
Copyright (c) 2016 Paul Masurel
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
https://github.com/quickwit-oss/bitpacking/blob/main/LICENSE
|
||||
@@ -0,0 +1,247 @@
|
||||
<div align="center">
|
||||
<p align="center">
|
||||
|
||||
<img width="257" alt="Lance Logo" src="https://user-images.githubusercontent.com/917119/199353423-d3e202f7-0269-411d-8ff2-e747e419e492.png">
|
||||
|
||||
**The Open Lakehouse Format for Multimodal AI**<br/>
|
||||
**High-performance vector search, full-text search, random access, and feature engineering capabilities for the lakehouse.**<br/>
|
||||
**Compatible with Pandas, DuckDB, Polars, PyArrow, Ray, Spark, and more integrations on the way.**
|
||||
|
||||
<a href="https://lance.org">Documentation</a> •
|
||||
<a href="https://lance.org/community">Community</a> •
|
||||
<a href="https://discord.gg/lance">Discord</a> •
|
||||
<a href="https://groups.google.com/a/lance.org/g/dev">Mailing List</a>
|
||||
|
||||
[CI]: https://github.com/lance-format/lance/actions/workflows/rust.yml
|
||||
[CI Badge]: https://github.com/lance-format/lance/actions/workflows/rust.yml/badge.svg
|
||||
[Docs]: https://lance.org
|
||||
[Docs Badge]: https://img.shields.io/badge/docs-passing-brightgreen
|
||||
[crates.io]: https://crates.io/crates/lance
|
||||
[crates.io badge]: https://img.shields.io/crates/v/lance.svg
|
||||
[Python versions]: https://pypi.org/project/pylance/
|
||||
[Python versions badge]: https://img.shields.io/pypi/pyversions/pylance
|
||||
|
||||
[![CI Badge]][CI]
|
||||
[![Docs Badge]][Docs]
|
||||
[![crates.io badge]][crates.io]
|
||||
[![Python versions badge]][Python versions]
|
||||
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<hr />
|
||||
|
||||
Lance is an open lakehouse format for multimodal AI. It contains a file format, table format, and catalog spec that allows you to build a complete lakehouse on top of object storage to power your AI workflows. Lance is perfect for:
|
||||
|
||||
1. Building search engines and feature stores with hybrid search capabilities.
|
||||
2. Large-scale ML training requiring high performance IO and random access.
|
||||
3. Storing, querying, and managing multimodal data including images, videos, audio, text, and embeddings.
|
||||
|
||||
The key features of Lance include:
|
||||
|
||||
* **Expressive hybrid search:** Combine vector similarity search, full-text search (BM25), and SQL analytics on the same dataset with accelerated secondary indices.
|
||||
|
||||
* **Lightning-fast random access:** 100x faster than Parquet or Iceberg for random access without sacrificing scan performance.
|
||||
|
||||
* **Native multimodal data support:** Store images, videos, audio, text, and embeddings in a single unified format with efficient blob encoding and lazy loading.
|
||||
|
||||
* **Data evolution:** Efficiently add columns with backfilled values without full table rewrites, perfect for ML feature engineering.
|
||||
|
||||
* **Zero-copy versioning:** Automatic versioning with ACID transactions, time travel, tags, and branches—no extra infrastructure needed.
|
||||
|
||||
* **Rich ecosystem integrations:** Apache Arrow, Pandas, Polars, DuckDB, Apache Spark, Ray, Trino, Apache Flink, and open catalogs (Apache Polaris, Unity Catalog, Apache Gravitino).
|
||||
|
||||
For more details, see the full [Lance format specification](https://lance.org/format).
|
||||
|
||||
> [!TIP]
|
||||
> Lance is in active development and we welcome contributions. Please see our [contributing guide](https://lance.org/community/contributing/) for more information.
|
||||
|
||||
## File format stability
|
||||
|
||||
Lance releases frequently because the SDKs, integrations, and performance work are moving quickly. This does not mean the Lance file format changes incompatibly in every release. The Lance file format is identified by the `data_storage_version` stored in each dataset, and stable storage versions are a long-term compatibility contract.
|
||||
|
||||
* Once a dataset is written with a stable `data_storage_version`, future Lance releases will continue to support reading that storage version.
|
||||
* SDK and API compatibility is separate from file format compatibility. SDK/API changes follow semantic versioning and are documented in the [migration guide](https://lance.org/guide/migration/).
|
||||
* Older Lance releases may not understand file format versions introduced later. If you run mixed Lance versions, pin `data_storage_version` for deterministic writes.
|
||||
* The `next` file format alias is unstable and should only be used for experimentation, never for production data.
|
||||
|
||||
For production, write data with a stable `data_storage_version`. See the [format versioning guide](https://lance.org/format/file/versioning/) for the current compatibility matrix.
|
||||
|
||||
## Quick Start
|
||||
|
||||
**Installation**
|
||||
|
||||
```shell
|
||||
pip install pylance
|
||||
```
|
||||
|
||||
To install a preview release:
|
||||
|
||||
```shell
|
||||
pip install --pre --extra-index-url https://pypi.fury.io/lance-format pylance
|
||||
```
|
||||
|
||||
> [!TIP]
|
||||
> Preview releases are released more often than full releases and contain the
|
||||
> latest features and bug fixes. They receive the same level of testing as full releases.
|
||||
> We guarantee they will remain published and available for download for at
|
||||
> least 6 months. When you want to pin to a specific version, prefer a stable release.
|
||||
|
||||
**Converting to Lance**
|
||||
|
||||
```python
|
||||
import lance
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pyarrow.dataset
|
||||
|
||||
df = pd.DataFrame({"a": [5], "b": [10]})
|
||||
uri = "/tmp/test.parquet"
|
||||
tbl = pa.Table.from_pandas(df)
|
||||
pa.dataset.write_dataset(tbl, uri, format='parquet')
|
||||
|
||||
parquet = pa.dataset.dataset(uri, format='parquet')
|
||||
lance.write_dataset(parquet, "/tmp/test.lance")
|
||||
```
|
||||
|
||||
**Reading Lance data**
|
||||
```python
|
||||
dataset = lance.dataset("/tmp/test.lance")
|
||||
assert isinstance(dataset, pa.dataset.Dataset)
|
||||
```
|
||||
|
||||
**Pandas**
|
||||
```python
|
||||
df = dataset.to_table().to_pandas()
|
||||
df
|
||||
```
|
||||
|
||||
**DuckDB**
|
||||
```python
|
||||
import duckdb
|
||||
|
||||
# If this segfaults, make sure you have duckdb v0.7+ installed
|
||||
duckdb.query("SELECT * FROM dataset LIMIT 10").to_df()
|
||||
```
|
||||
|
||||
**Vector search**
|
||||
|
||||
Download the sift1m subset
|
||||
|
||||
```shell
|
||||
wget ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz
|
||||
tar -xzf sift.tar.gz
|
||||
```
|
||||
|
||||
Convert it to Lance
|
||||
|
||||
```python
|
||||
import lance
|
||||
from lance.vector import vec_to_table
|
||||
import numpy as np
|
||||
import struct
|
||||
|
||||
nvecs = 1000000
|
||||
ndims = 128
|
||||
with open("sift/sift_base.fvecs", mode="rb") as fobj:
|
||||
buf = fobj.read()
|
||||
data = np.array(struct.unpack("<128000000f", buf[4 : 4 + 4 * nvecs * ndims])).reshape((nvecs, ndims))
|
||||
dd = dict(zip(range(nvecs), data))
|
||||
|
||||
table = vec_to_table(dd)
|
||||
uri = "vec_data.lance"
|
||||
sift1m = lance.write_dataset(table, uri, max_rows_per_group=8192, max_rows_per_file=1024*1024)
|
||||
```
|
||||
|
||||
Build the index
|
||||
|
||||
```python
|
||||
sift1m.create_index("vector",
|
||||
index_type="IVF_PQ",
|
||||
num_partitions=256, # IVF
|
||||
num_sub_vectors=16) # PQ
|
||||
```
|
||||
|
||||
Search the dataset
|
||||
|
||||
```python
|
||||
# Get top 10 similar vectors
|
||||
import duckdb
|
||||
|
||||
dataset = lance.dataset(uri)
|
||||
|
||||
# Sample 100 query vectors. If this segfaults, make sure you have duckdb v0.7+ installed
|
||||
sample = duckdb.query("SELECT vector FROM dataset USING SAMPLE 100").to_df()
|
||||
query_vectors = np.array([np.array(x) for x in sample.vector])
|
||||
|
||||
# Get nearest neighbors for all of them
|
||||
rs = [dataset.to_table(nearest={"column": "vector", "k": 10, "q": q})
|
||||
for q in query_vectors]
|
||||
```
|
||||
|
||||
## Directory structure
|
||||
|
||||
| Directory | Description |
|
||||
|--------------------|--------------------------|
|
||||
| [rust](./rust) | Core Rust implementation |
|
||||
| [python](./python) | Python bindings (PyO3) |
|
||||
| [java](./java) | Java bindings (JNI) |
|
||||
| [docs](./docs) | Documentation source |
|
||||
|
||||
## Benchmarks
|
||||
|
||||
### Vector search
|
||||
|
||||
We used the SIFT dataset to benchmark our results with 1M vectors of 128D
|
||||
|
||||
1. For 100 randomly sampled query vectors, we get <1ms average response time (on a 2023 m2 MacBook Air)
|
||||
|
||||

|
||||
|
||||
2. ANNs are always a trade-off between recall and performance
|
||||
|
||||

|
||||
|
||||
### Vs. parquet
|
||||
|
||||
We create a Lance dataset using the Oxford Pet dataset to do some preliminary performance testing of Lance as compared to Parquet and raw image/XMLs. For analytics queries, Lance is 50-100x better than reading the raw metadata. For batched random access, Lance is 100x better than both parquet and raw files.
|
||||
|
||||

|
||||
|
||||
## Why Lance for AI/ML workflows?
|
||||
|
||||
The machine learning development cycle involves multiple stages:
|
||||
|
||||
```mermaid
|
||||
graph LR
|
||||
A[Collection] --> B[Exploration];
|
||||
B --> C[Analytics];
|
||||
C --> D[Feature Engineer];
|
||||
D --> E[Training];
|
||||
E --> F[Evaluation];
|
||||
F --> C;
|
||||
E --> G[Deployment];
|
||||
G --> H[Monitoring];
|
||||
H --> A;
|
||||
```
|
||||
|
||||
Traditional lakehouse formats were designed for SQL analytics and struggle with AI/ML workloads that require:
|
||||
- **Vector search** for similarity and semantic retrieval
|
||||
- **Fast random access** for sampling and interactive exploration
|
||||
- **Multimodal data** storage (images, videos, audio alongside embeddings)
|
||||
- **Data evolution** for feature engineering without full table rewrites
|
||||
- **Hybrid search** combining vectors, full-text, and SQL predicates
|
||||
|
||||
While existing formats (Parquet, Iceberg, Delta Lake) excel at SQL analytics, they require additional specialized systems for AI capabilities. Lance brings these AI-first features directly into the lakehouse format.
|
||||
|
||||
A comparison of different formats across ML development stages:
|
||||
|
||||
| | Lance | Parquet & ORC | JSON & XML | TFRecord | Database | Warehouse |
|
||||
|---------------------|-------|---------------|------------|----------|----------|-----------|
|
||||
| Analytics | Fast | Fast | Slow | Slow | Decent | Fast |
|
||||
| Feature Engineering | Fast | Fast | Decent | Slow | Decent | Good |
|
||||
| Training | Fast | Decent | Slow | Fast | N/A | N/A |
|
||||
| Exploration | Fast | Slow | Fast | Slow | Fast | Decent |
|
||||
| Infra Support | Rich | Rich | Decent | Limited | Rich | Rich |
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
# Protobuf Guidelines
|
||||
|
||||
Also see [root AGENTS.md](../AGENTS.md) for cross-language standards.
|
||||
|
||||
## Compatibility
|
||||
|
||||
- Protobuf schemas that are part of a stable file format or any other stable persisted contract must remain backwards compatible. Never reuse or change their existing field numbers.
|
||||
- Protobuf schemas used exclusively by an unstable file format follow the root file-format stability contract: do not preserve compatibility with prior unstable revisions. Before making a breaking protobuf change, verify that the schema is not shared with a stable format or another persisted contract.
|
||||
|
||||
## Schema Design
|
||||
|
||||
- Use `optional` when you need to distinguish "not set" from "zero value" — `optional` enables presence tracking (`has_*` methods) and maps to `Option<T>` in Rust. Bare proto3 fields have no presence semantics: they always hold a value (defaulting to zero), so you cannot tell if the sender explicitly set them.
|
||||
- Use structured message types (e.g., `BasePath`) instead of plain scalars, and scope fields to operation-specific messages (e.g., `InsertTransaction`) rather than generic top-level ones.
|
||||
- Don't duplicate data across messages — store each fact once and derive relationships. Prefer parallel sequences over maps when keys already exist in another field.
|
||||
|
||||
## Documentation
|
||||
|
||||
- Document the semantic meaning of both present and absent states for `optional` fields — explain when each case applies.
|
||||
- Use precise domain terminology in field descriptions — avoid ambiguous abbreviations or terms that collide with domain concepts.
|
||||
Symlink
+1
@@ -0,0 +1 @@
|
||||
AGENTS.md
|
||||
@@ -0,0 +1,72 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.pb;
|
||||
|
||||
import "table_identifier.proto";
|
||||
import "table.proto";
|
||||
import "index.proto";
|
||||
|
||||
// Query-time approximation mode for vector search.
|
||||
//
|
||||
// This currently only affects RQ-quantized vector indexes, such as IVF_RQ.
|
||||
// Other index types ignore this setting.
|
||||
enum VectorApproxMode {
|
||||
// Use all RQ bits for query-time scoring with u8-quantized lookup tables.
|
||||
Normal = 0;
|
||||
// Use only one RQ bit for query-time scoring, even for multi-bit indexes.
|
||||
Fast = 1;
|
||||
// Use all RQ bits for query-time scoring with u16-quantized lookup tables
|
||||
// to reduce estimator quantization error.
|
||||
Accurate = 2;
|
||||
}
|
||||
|
||||
// Serialized vector query parameters.
|
||||
message VectorQueryProto {
|
||||
// Query vector as Arrow IPC bytes (supports Float16, Float32, Float64, UInt8, etc.)
|
||||
bytes query_vector_arrow_ipc = 1;
|
||||
string column = 2;
|
||||
uint32 k = 3;
|
||||
optional float lower_bound = 4;
|
||||
optional float upper_bound = 5;
|
||||
optional uint32 minimum_nprobes = 6;
|
||||
optional uint32 maximum_nprobes = 7;
|
||||
optional uint32 ef = 8;
|
||||
optional uint32 refine_factor = 9;
|
||||
// Distance metric type. Absent means None (use the index's default metric).
|
||||
optional lance.index.pb.VectorMetricType metric_type = 10;
|
||||
bool use_index = 11;
|
||||
optional float dist_q_c = 12;
|
||||
optional int32 query_parallelism = 13;
|
||||
// Query-time approximation mode. Currently only affects RQ-quantized vector
|
||||
// indexes, such as IVF_RQ. Other index types ignore this setting.
|
||||
VectorApproxMode approx_mode = 14;
|
||||
}
|
||||
|
||||
// Serializable form of ANNIvfSubIndexExec — the IVF sub-index search node.
|
||||
//
|
||||
// The prefilter child ExecutionPlan is serialized by DataFusion's codec
|
||||
// automatically via children() / with_new_children(). The prefilter_type
|
||||
// field tells the decoder which PreFilterSource variant to use when
|
||||
// reconstructing from the deserialized child inputs.
|
||||
message ANNIvfSubIndexExecProto {
|
||||
enum PreFilterType {
|
||||
NONE = 0;
|
||||
FILTERED_ROW_IDS = 1;
|
||||
SCALAR_INDEX_QUERY = 2;
|
||||
}
|
||||
|
||||
VectorQueryProto query = 1;
|
||||
lance.datafusion.TableIdentifier table = 2;
|
||||
repeated lance.table.IndexMetadata indices = 3;
|
||||
PreFilterType prefilter_type = 4;
|
||||
}
|
||||
|
||||
// Serializable form of ANNIvfPartitionExec — the IVF centroid routing node.
|
||||
message ANNIvfPartitionExecProto {
|
||||
VectorQueryProto query = 1;
|
||||
lance.datafusion.TableIdentifier table = 2;
|
||||
repeated string index_uuids = 3;
|
||||
}
|
||||
@@ -0,0 +1,347 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.encodings;
|
||||
|
||||
import "google/protobuf/empty.proto";
|
||||
|
||||
// This file contains a specification for encodings that can be used
|
||||
// to store and load Arrow data into a Lance file for the 2.0 format. It
|
||||
// has been superseded by encodings21.proto which is used for the 2.1 format.
|
||||
//
|
||||
// # Types
|
||||
//
|
||||
// This file assumes the user wants to load data into Arrow arrays and
|
||||
// explains how to map Arrow arrays into Lance files. Encodings are divided
|
||||
// into "array encoding" (which maps to an Arrow array and may contain multiple
|
||||
// buffers) and "buffer encoding" (which encodes a single buffer of data).
|
||||
//
|
||||
// # Encoding Tree
|
||||
//
|
||||
// Most encodings are layered on top of each other. These form a tree of
|
||||
// encodings with a single root node. To encode an array you will typically
|
||||
// start with the root node and then take the output from that root encoding
|
||||
// and feed it into child encodings. The decoding process works in reverse.
|
||||
//
|
||||
// # Multi-column Encodings
|
||||
//
|
||||
// Some Arrow arrays will map to more than one column of Lance data. For
|
||||
// example, struct arrays and list arrays. This file only contains encodings
|
||||
// for a single column. However, it does describe how multi-column arrays can
|
||||
// be encoded.
|
||||
|
||||
// A pointer to a buffer in a Lance file
|
||||
//
|
||||
// A writer can place a buffer in three different locations. The buffer
|
||||
// can go in the data page, in the column metadata, or in the file metadata.
|
||||
// The writer is free to choose whatever is most appropriate (for example, a dictionary
|
||||
// that is shared across all pages in a column will probably go in the column
|
||||
// metadata). This specification does not dictate where the buffer should go.
|
||||
message Buffer {
|
||||
// The index of the buffer in the collection of buffers
|
||||
uint32 buffer_index = 1;
|
||||
// The collection holding the buffer
|
||||
enum BufferType {
|
||||
// The buffer is stored in the data page itself
|
||||
page = 0;
|
||||
// The buffer is stored in the column metadata
|
||||
column = 1;
|
||||
// The buffer is stored in the file metadata
|
||||
file = 2;
|
||||
};
|
||||
BufferType buffer_type = 2;
|
||||
}
|
||||
|
||||
// An encoding that adds nullability to another array encoding
|
||||
//
|
||||
// This can wrap any array encoding and add nullability information
|
||||
message Nullable {
|
||||
message NoNull {
|
||||
ArrayEncoding values = 1;
|
||||
}
|
||||
message AllNull {}
|
||||
message SomeNull {
|
||||
ArrayEncoding validity = 1;
|
||||
ArrayEncoding values = 2;
|
||||
}
|
||||
oneof nullability {
|
||||
// The array has no nulls and there is a single buffer needed
|
||||
NoNull no_nulls = 1;
|
||||
// The array may have nulls and we need two buffers
|
||||
SomeNull some_nulls = 2;
|
||||
// All values are null (no buffers needed)
|
||||
AllNull all_nulls = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// An array encoding for variable-length list fields
|
||||
message List {
|
||||
// An array containing the offsets into an items array.
|
||||
//
|
||||
// This array will have num_rows items and will never
|
||||
// have nulls.
|
||||
//
|
||||
// If the list at index i is not null then offsets[i] will
|
||||
// contain `base + len(list)` where `base` is defined as:
|
||||
// i == 0: 0
|
||||
// i > 0: (offsets[i-1] % null_offset_adjustment)
|
||||
//
|
||||
// To help understand we can consider the following example list:
|
||||
// [ [A, B], null, [], [C, D, E] ]
|
||||
//
|
||||
// The offsets will be [2, ?, 2, 5]
|
||||
//
|
||||
// If the incoming list at index i IS null then offsets[i] will
|
||||
// contain `base + len(list) + null_offset_adjustment` where `base`
|
||||
// is defined the same as above.
|
||||
//
|
||||
// To complete the above example let's assume that `null_offset_adjustment`
|
||||
// is 7. Then the offsets will be [2, 9, 2, 5]
|
||||
//
|
||||
// If there are no nulls then the offsets we write here are exactly the
|
||||
// same as the offsets in an Arrow list array (except we omit the leading
|
||||
// 0 which is redundant)
|
||||
//
|
||||
// The reason we do this is so that reading a single list at index i only
|
||||
// requires us to load the indices at i and i-1.
|
||||
//
|
||||
// If the offset at index i is greater than `null_offset_adjustment``
|
||||
// then the list at index i is null.
|
||||
//
|
||||
// Otherwise the length of the list is `offsets[i] - base` where
|
||||
// base is defined the same as above.
|
||||
//
|
||||
// Let's consider our example offsets: [2, 9, 2, 5]
|
||||
//
|
||||
// We can take any range of lists and determine how many list items are
|
||||
// referenced by the sublist.
|
||||
//
|
||||
// 0..3: [_, 5] -> items 0..5 (base = 0* and end is 5)
|
||||
// 0..2: [_, 2] -> items 0..2 (base = 0* and end is 2)
|
||||
// 0..1: [_, 9] -> items 0..2 (base = 0* and end is 9 % 7)
|
||||
// 1..3: [2, 5] -> items 2..5 (base = 2 and end is 5)
|
||||
// 1..2: [2, 2] -> items 2..2 (base = 2 and end is 2)
|
||||
// 2..3: [9, 5] -> items 2..5 (base = 9 % 7 and end is 5)
|
||||
//
|
||||
// * When the start of our range is the 0th item the base is always 0 and we only
|
||||
// need to load a single index from disk to determine the range.
|
||||
//
|
||||
// The data type of the offsets array is flexible and does not need
|
||||
// to match the data type of the destination array. Please note that the offsets
|
||||
// array is very likely to be efficiently encoded by bit packing deltas.
|
||||
ArrayEncoding offsets = 1;
|
||||
// If a list is null then we add this value to the offset
|
||||
//
|
||||
// This value must be greater than the length of the items so that
|
||||
// (offset + null_offset_adjustment) is never used by a non-null list.
|
||||
//
|
||||
// Note that this value cannot be equal to the length of the items
|
||||
// because then a page with a single list would store [ X ] and we
|
||||
// couldn't know if that is a null list or a list with X items.
|
||||
//
|
||||
// Therefore, the best choice for this value is 1 + # of items.
|
||||
// Choosing this will maximize the bit packing that we can apply to the offsets.
|
||||
uint64 null_offset_adjustment = 2;
|
||||
// How many items are referenced by these offsets. This is needed in
|
||||
// order to determine which items pages map to this offsets page.
|
||||
uint64 num_items = 3;
|
||||
}
|
||||
|
||||
// An array encoding for fixed-size list fields
|
||||
message FixedSizeList {
|
||||
/// The number of items in each list
|
||||
uint32 dimension = 1;
|
||||
/// True if the list is nullable
|
||||
bool has_validity = 3;
|
||||
/// The items in the list
|
||||
ArrayEncoding items = 2;
|
||||
}
|
||||
|
||||
message Compression {
|
||||
string scheme = 1;
|
||||
optional int32 level = 2;
|
||||
}
|
||||
|
||||
// Fixed width items placed contiguously in a buffer
|
||||
message Flat {
|
||||
// the number of bits per value, must be greater than 0, does
|
||||
// not need to be a multiple of 8
|
||||
uint64 bits_per_value = 1;
|
||||
// the buffer of values
|
||||
Buffer buffer = 2;
|
||||
// The Compression message can specify the compression scheme (e.g. zstd) and any
|
||||
// other information that is needed for decompression.
|
||||
//
|
||||
// If this array is compressed then the bits_per_value refers to the uncompressed
|
||||
// data.
|
||||
Compression compression = 3;
|
||||
}
|
||||
|
||||
// Compression algorithm where all values have a constant value
|
||||
message Constant {
|
||||
// The value (TODO: define encoding for literals?)
|
||||
bytes value = 1;
|
||||
}
|
||||
|
||||
// Items are bitpacked in a buffer
|
||||
message Bitpacked {
|
||||
// the number of bits used for a value in the buffer
|
||||
uint64 compressed_bits_per_value = 1;
|
||||
|
||||
// the number of bits of the uncompressed value. e.g. for a u32, this will be 32
|
||||
uint64 uncompressed_bits_per_value = 2;
|
||||
|
||||
// The items in the list
|
||||
Buffer buffer = 3;
|
||||
|
||||
// Whether or not a sign bit is included in the bitpacked value
|
||||
bool signed = 4;
|
||||
}
|
||||
|
||||
// Items are bitpacked in a buffer
|
||||
message BitpackedForNonNeg {
|
||||
// the number of bits used for a value in the buffer
|
||||
uint64 compressed_bits_per_value = 1;
|
||||
|
||||
// the number of bits of the uncompressed value. e.g. for a u32, this will be 32
|
||||
uint64 uncompressed_bits_per_value = 2;
|
||||
|
||||
// The items in the list
|
||||
Buffer buffer = 3;
|
||||
}
|
||||
|
||||
// Opaque bitpacking variant where the bits per value are stored inline in the chunks themselves
|
||||
message InlineBitpacking {
|
||||
// the number of bits of the uncompressed value. e.g. for a u32, this will be 32
|
||||
uint64 uncompressed_bits_per_value = 2;
|
||||
}
|
||||
|
||||
// Transparent bitpacking variant where the number of bits per value is fixed through the whole buffer
|
||||
message OutOfLineBitpacking {
|
||||
// the number of bits of the uncompressed value. e.g. for a u32, this will be 32
|
||||
uint64 uncompressed_bits_per_value = 2;
|
||||
// The number of compressed bits per value, fixed across the entire buffer
|
||||
uint64 compressed_bits_per_value = 3;
|
||||
}
|
||||
|
||||
// An array encoding for shredded structs that will never be null
|
||||
//
|
||||
// There is no actual data in this column.
|
||||
//
|
||||
// TODO: Struct validity bitmaps will be placed here.
|
||||
message SimpleStruct {}
|
||||
|
||||
// An array encoding for binary fields
|
||||
message Binary {
|
||||
ArrayEncoding indices = 1;
|
||||
ArrayEncoding bytes = 2;
|
||||
uint64 null_adjustment = 3;
|
||||
}
|
||||
|
||||
message Variable {
|
||||
uint32 bits_per_offset = 1;
|
||||
}
|
||||
|
||||
message Fsst {
|
||||
ArrayEncoding binary = 1;
|
||||
bytes symbol_table = 2;
|
||||
}
|
||||
|
||||
// An array encoding for dictionary-encoded fields
|
||||
message Dictionary {
|
||||
ArrayEncoding indices = 1;
|
||||
ArrayEncoding items = 2;
|
||||
uint32 num_dictionary_items = 3;
|
||||
}
|
||||
|
||||
message PackedStruct {
|
||||
repeated ArrayEncoding inner = 1;
|
||||
Buffer buffer = 2;
|
||||
}
|
||||
|
||||
message PackedStructFixedWidthMiniBlock {
|
||||
ArrayEncoding Flat = 1;
|
||||
repeated uint32 bits_per_values = 2;
|
||||
}
|
||||
|
||||
message FixedSizeBinary {
|
||||
ArrayEncoding bytes = 1;
|
||||
uint32 byte_width = 2;
|
||||
}
|
||||
|
||||
message Block {
|
||||
string scheme = 1;
|
||||
}
|
||||
|
||||
// Run-Length Encoding for miniblock format
|
||||
message Rle {
|
||||
// Number of bits per value (8, 16, 32, 64, or 128)
|
||||
uint64 bits_per_value = 1;
|
||||
}
|
||||
|
||||
// Byte Stream Split encoding for floating point values
|
||||
message ByteStreamSplit {
|
||||
// Number of bits per value (32 for float, 64 for double)
|
||||
uint64 bits_per_value = 1;
|
||||
}
|
||||
|
||||
// General miniblock encoding - wraps another miniblock encoding with compression
|
||||
message GeneralMiniBlock {
|
||||
// The inner miniblock encoding (e.g., Rle, Bitpacked, etc.)
|
||||
ArrayEncoding inner = 1;
|
||||
// The compression scheme to apply to the miniblock buffers
|
||||
Compression compression = 2;
|
||||
}
|
||||
|
||||
// Encodings that decode into an Arrow array
|
||||
message ArrayEncoding {
|
||||
oneof array_encoding {
|
||||
Flat flat = 1;
|
||||
Nullable nullable = 2;
|
||||
FixedSizeList fixed_size_list = 3;
|
||||
List list = 4;
|
||||
SimpleStruct struct = 5;
|
||||
Binary binary = 6;
|
||||
Dictionary dictionary = 7;
|
||||
Fsst fsst = 8;
|
||||
PackedStruct packed_struct = 9;
|
||||
Bitpacked bitpacked = 10;
|
||||
FixedSizeBinary fixed_size_binary = 11;
|
||||
BitpackedForNonNeg bitpacked_for_non_neg = 12;
|
||||
Constant constant = 13;
|
||||
InlineBitpacking inline_bitpacking = 14;
|
||||
OutOfLineBitpacking out_of_line_bitpacking = 15;
|
||||
Variable variable = 16;
|
||||
PackedStructFixedWidthMiniBlock packed_struct_fixed_width_mini_block = 17;
|
||||
Block block = 18;
|
||||
Rle rle = 19;
|
||||
GeneralMiniBlock general_mini_block = 20;
|
||||
ByteStreamSplit byte_stream_split = 21;
|
||||
}
|
||||
}
|
||||
|
||||
// Wraps a column with a zone map index that can be used
|
||||
// to apply pushdown filters
|
||||
message ZoneIndex {
|
||||
uint32 rows_per_zone = 1;
|
||||
Buffer zone_map_buffer = 2;
|
||||
ColumnEncoding inner = 3;
|
||||
}
|
||||
|
||||
// Marks a column as blob data. It will contain a packed struct
|
||||
// with fields position and size (u64)
|
||||
message Blob {
|
||||
ColumnEncoding inner = 1;
|
||||
}
|
||||
|
||||
// Encodings that describe a column of values
|
||||
message ColumnEncoding {
|
||||
oneof column_encoding {
|
||||
// No special encoding, just column values
|
||||
google.protobuf.Empty values = 1;
|
||||
ZoneIndex zone_index = 2;
|
||||
Blob blob = 3;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,635 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.encodings21;
|
||||
|
||||
// This file contains a specification for encodings that can be used
|
||||
// to store and load Arrow data into a Lance file for the 2.1 format.
|
||||
//
|
||||
// # Types
|
||||
//
|
||||
// This file assumes the user wants to load data into Arrow arrays and
|
||||
// explains how to map Arrow arrays into Lance files. Encodings are divided
|
||||
// into "structural encodings" (which are used to encode the structure of the
|
||||
// data such as any list or struct layers) and "compressive encodings" (which
|
||||
// are used to compress the actual data values).
|
||||
//
|
||||
// # Standardized Interpretation of Counting Terms
|
||||
//
|
||||
// When working with 2.1 encodings we have a number of different "counting terms" and it can be
|
||||
// difficult to understand what we mean when we are talking about a "number of values". Here is
|
||||
// a standard interpretation of these terms:
|
||||
//
|
||||
// To understand these definitions consider a data type FIXED_SIZE_LIST<LIST<INT32>>.
|
||||
//
|
||||
// A "value" is an abstract term when we aren't being specific.
|
||||
//
|
||||
// - num_rows: This is the highest level counting term. A single row includes everything in the
|
||||
// fixed size list. This is what the user asks for when they asks for a range of rows.
|
||||
// - num_elements: The number of elements is the number of rows multiplied by the dimension of any
|
||||
// fixed size list wrappers. This is what you get when you flatten the FSL layer and
|
||||
// is the starting point for structural encoding. Note that an element can be a list
|
||||
// value or a single primitive value.
|
||||
// - num_items: The number of items is the number of values in the repetition and definition vectors
|
||||
// after everything has been flattened.
|
||||
// - num_visible_items: The number of visible items is the number of items after invisible items
|
||||
// have been removed. Invisible items are rep/def levels that don't correspond to an
|
||||
// actual value.
|
||||
|
||||
|
||||
// # Structural Encodings
|
||||
//
|
||||
// The following message are used to describe the structural encoding of the
|
||||
// data. In this document, we refer to these structural encodings as layouts.
|
||||
|
||||
// Repetition and definition levels are described in more detail elsewhere. As we peel through
|
||||
// the structure of an array we will encounter layers of struct and list. Each of these layers
|
||||
// potentially adds a new level to the repetition and definition levels. This message describes
|
||||
// the meaning of each layer.
|
||||
enum RepDefLayer {
|
||||
// Should never be used, included for debugging purporses and general protobuf best practice
|
||||
REPDEF_UNSPECIFIED = 0;
|
||||
// All values are valid (can be primitive or struct)
|
||||
REPDEF_ALL_VALID_ITEM = 1;
|
||||
// All list values are valid
|
||||
REPDEF_ALL_VALID_LIST = 2;
|
||||
// There are one or more null items (can be primitive or struct)
|
||||
REPDEF_NULLABLE_ITEM = 3;
|
||||
// A list layer with null lists but no empty lists
|
||||
REPDEF_NULLABLE_LIST = 4;
|
||||
// A list layer with empty lists but no null lists
|
||||
REPDEF_EMPTYABLE_LIST = 5;
|
||||
// A list layer with both empty lists and null lists
|
||||
REPDEF_NULL_AND_EMPTY_LIST = 6;
|
||||
}
|
||||
|
||||
// A layout used for pages where the data is small
|
||||
//
|
||||
// In this case we can fit many values into a single disk sector and transposing buffers is
|
||||
// expensive. As a result, we do not transpose the buffers but compress the data into small
|
||||
// chunks (called mini blocks) which are roughly the size of a disk sector.
|
||||
//
|
||||
// The end result is a small amount of read amplification (since we must read an entire page
|
||||
// at a time) but we have more flexibility in compression and do less work per value when
|
||||
// compressing and decompressing in bulk.
|
||||
message MiniBlockLayout {
|
||||
// Description of the compression of repetition levels (e.g. how many bits per rep)
|
||||
//
|
||||
// Optional, if there is no repetition then this field is not present
|
||||
CompressiveEncoding rep_compression = 1;
|
||||
// Description of the compression of definition levels (e.g. how many bits per def)
|
||||
//
|
||||
// Optional, if there is no definition then this field is not present
|
||||
CompressiveEncoding def_compression = 2;
|
||||
// Description of the compression of values
|
||||
CompressiveEncoding value_compression = 3;
|
||||
// Description of the compression of the dictionary data
|
||||
//
|
||||
// Optional, if there is no dictionary then this field is not present
|
||||
CompressiveEncoding dictionary = 4;
|
||||
// Number of items in the dictionary
|
||||
uint64 num_dictionary_items = 5;
|
||||
// The meaning of each repdef layer, used to interpret repdef buffers correctly
|
||||
repeated RepDefLayer layers = 6;
|
||||
// The number of buffers in each mini-block, this is determined by the compression and does
|
||||
// NOT include the repetition or definition buffers (the presence of these buffers can be determined
|
||||
// by looking at the rep_compression and def_compression fields)
|
||||
uint64 num_buffers = 7;
|
||||
// The depth of the repetition index.
|
||||
//
|
||||
// If there is repetition then the depth must be at least 1. If there are many layers
|
||||
// of repetition then deeper repetition indices will support deeper nested random access. For
|
||||
// example, given 5 layers of repetition then the repetition index depth must be at least
|
||||
// 3 to support access like `rows[50][17][3]`.
|
||||
//
|
||||
// We require `repetition_index_depth + 1` u64 values per mini-block to store the repetition
|
||||
// index if the `repetition_index_depth` is greater than 0. The +1 is because we need to store
|
||||
// the number of "leftover items" at the end of the chunk. Otherwise, we wouldn't have any way
|
||||
// to know if the final item in a chunk is valid or not.
|
||||
uint32 repetition_index_depth = 8;
|
||||
// The page already records how many rows are in the page. For mini-block we also need to know how
|
||||
// many "items" are in the page. A row and an item are the same thing unless the page has lists.
|
||||
uint64 num_items = 9;
|
||||
|
||||
// Since Lance 2.2, miniblocks have larger chunk sizes (>= 64KB)
|
||||
bool has_large_chunk = 10;
|
||||
}
|
||||
|
||||
// A layout used for pages where the data is large
|
||||
//
|
||||
// In this case the cost of transposing the data is relatively small (compared to the cost of writing the data)
|
||||
// and so we just zip the buffers together
|
||||
message FullZipLayout {
|
||||
// The number of bits of repetition info (0 if there is no repetition)
|
||||
uint32 bits_rep = 1;
|
||||
// The number of bits of definition info (0 if there is no definition)
|
||||
uint32 bits_def = 2;
|
||||
// The number of bits of value info
|
||||
//
|
||||
// Note: we use bits here (and not bytes) for consistency with other encodings. However, in practice,
|
||||
// there is never a reason to use a bits per value that is not a multiple of 8. The complexity is not
|
||||
// worth the small savings in space since this encoding is typically used with large values already.
|
||||
oneof details {
|
||||
// If this is a fixed width block then we need to have a fixed number of bits per value
|
||||
uint32 bits_per_value = 3;
|
||||
// If this is a variable width block then we need to have a fixed number of bits per offset
|
||||
uint32 bits_per_offset = 4;
|
||||
}
|
||||
// The number of items in the page
|
||||
uint32 num_items = 5;
|
||||
// The number of visible items in the page
|
||||
uint32 num_visible_items = 6;
|
||||
// Description of the compression of values
|
||||
CompressiveEncoding value_compression = 7;
|
||||
// The meaning of each repdef layer, used to interpret repdef buffers correctly
|
||||
repeated RepDefLayer layers = 8;
|
||||
}
|
||||
|
||||
// A layout used for sparse flat or nested pages where Arrow structure is represented directly
|
||||
// in layer-local slot domains instead of as dense repetition / definition events.
|
||||
//
|
||||
// Structural layers are ordered from outer-most to inner-most. Values remain mini-block
|
||||
// compressed and are split into independently readable chunks.
|
||||
message SparseLayout {
|
||||
// Description of the compression of values.
|
||||
CompressiveEncoding value_compression = 1;
|
||||
// Number of value buffers in each mini-block chunk. This does not include structural buffers.
|
||||
uint64 num_buffers = 2;
|
||||
// Number of entries in the equivalent dense repetition / definition stream. This equals
|
||||
// num_visible_items plus one structural placeholder for every list slot without children.
|
||||
// Null leaf slots count as visible items because they still occupy positions in Arrow's
|
||||
// leaf value buffer. For example, a nullable primitive with 100 slots, 30 of them null,
|
||||
// has num_items = num_visible_items = 100.
|
||||
uint64 num_items = 3;
|
||||
// Number of leaf value slots encoded in the value chunks, including null leaf slots.
|
||||
uint64 num_visible_items = 4;
|
||||
// If true, chunk-local value buffer sizes use u32. Otherwise they use u16.
|
||||
bool has_large_chunk = 5;
|
||||
// Structural layers ordered from outer-most to inner-most. This may be empty for a flat,
|
||||
// non-nullable leaf page whose scheduling domain equals num_visible_items.
|
||||
repeated SparseStructuralLayer structural_layers = 6;
|
||||
}
|
||||
|
||||
// A domain is a layer-local integer coordinate space [0, num_slots). A slot is one
|
||||
// element in that space. The outer-most domain is the page's top-level rows; each
|
||||
// layer's child domain is the next layer's parent domain, and the terminal child
|
||||
// domain contains num_visible_items leaf value slots.
|
||||
message SparseStructuralLayer {
|
||||
// Exactly one layer kind is required.
|
||||
oneof layer {
|
||||
SparseValidityLayer validity = 1;
|
||||
SparseListLayer list = 2;
|
||||
SparseFixedSizeListLayer fixed_size_list = 3;
|
||||
}
|
||||
}
|
||||
|
||||
message SparseValidityLayer {
|
||||
// Number of nullable item or struct slots in this layer's parent and child domain.
|
||||
uint64 num_slots = 1;
|
||||
// Validity for the slots in this layer.
|
||||
SparseValiditySet validity = 2;
|
||||
}
|
||||
|
||||
message SparseListLayer {
|
||||
// Number of list, large-list, or map slots in this layer's parent domain.
|
||||
uint64 num_slots = 1;
|
||||
// Number of slots in this layer's child domain.
|
||||
uint64 num_child_slots = 2;
|
||||
// Non-empty parent slots. Valid parent slots absent from this set are empty lists.
|
||||
SparsePositionSet non_empty_positions = 3;
|
||||
// Positive child counts corresponding one-for-one with non_empty_positions.
|
||||
SparseCountSet counts = 4;
|
||||
// Validity for the parent slots in this layer.
|
||||
SparseValiditySet validity = 5;
|
||||
}
|
||||
|
||||
message SparseFixedSizeListLayer {
|
||||
// Number of fixed-size-list slots in this layer's parent domain.
|
||||
uint64 num_slots = 1;
|
||||
// Number of children per parent slot. The child domain has num_slots * dimension slots.
|
||||
uint64 dimension = 2;
|
||||
// Validity for the parent slots in this layer.
|
||||
SparseValiditySet validity = 3;
|
||||
}
|
||||
|
||||
message SparseValiditySet {
|
||||
enum Meaning {
|
||||
SPARSE_VALIDITY_UNSPECIFIED = 0;
|
||||
// Stored positions are null; all other positions are valid.
|
||||
SPARSE_VALIDITY_NULL_POSITIONS = 1;
|
||||
// Stored positions are valid; all other positions are null.
|
||||
SPARSE_VALIDITY_VALID_POSITIONS = 2;
|
||||
}
|
||||
|
||||
Meaning meaning = 1;
|
||||
SparsePositionSet positions = 2;
|
||||
}
|
||||
|
||||
message SparsePositionEmpty {}
|
||||
|
||||
message SparsePositionAll {}
|
||||
|
||||
message SparsePositionRange {
|
||||
uint64 start = 1;
|
||||
uint64 length = 2;
|
||||
}
|
||||
|
||||
message SparsePositionSet {
|
||||
oneof positions {
|
||||
// Delta-compressed u64 positions. Cardinality is num_positions.
|
||||
CompressiveEncoding explicit = 1;
|
||||
// One contiguous, non-empty range.
|
||||
SparsePositionRange range = 2;
|
||||
// Every position in the domain.
|
||||
SparsePositionAll all = 3;
|
||||
// No positions in the domain.
|
||||
SparsePositionEmpty empty = 4;
|
||||
}
|
||||
// Semantic cardinality of this set.
|
||||
uint64 num_positions = 5;
|
||||
}
|
||||
|
||||
message SparseCountEmpty {}
|
||||
|
||||
message SparseCountConstant {
|
||||
// Child count shared by every non-empty list slot.
|
||||
uint64 value = 1;
|
||||
}
|
||||
|
||||
message SparseCountSet {
|
||||
oneof counts {
|
||||
// Compressed u64 child counts. Cardinality comes from the containing position set.
|
||||
CompressiveEncoding explicit = 1;
|
||||
// One positive child count shared by every non-empty list slot.
|
||||
SparseCountConstant constant = 2;
|
||||
// No counts; valid only when there are no non-empty list slots.
|
||||
SparseCountEmpty empty = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// A layout used for pages where all (visible) values are the same scalar value.
|
||||
//
|
||||
// This generalizes the prior AllNullLayout semantics for file_version >= 2.2.
|
||||
//
|
||||
// There may be buffers of repetition and definition information if required in order
|
||||
// to interpret what kind of nulls are present / which items are visible.
|
||||
message ConstantLayout {
|
||||
// The meaning of each repdef layer, used to interpret repdef buffers correctly
|
||||
repeated RepDefLayer layers = 5;
|
||||
|
||||
// Inline fixed-width scalar value bytes.
|
||||
//
|
||||
// This MUST only be used for types where a single non-null element is represented by a single
|
||||
// fixed-width Arrow value buffer (i.e. no offsets buffer, no child data).
|
||||
//
|
||||
// Constraints:
|
||||
// - MUST be absent for an all-null page
|
||||
// - MUST be <= 32 bytes if present
|
||||
optional bytes inline_value = 6;
|
||||
|
||||
// Optional compression algorithm used for the repetition buffer.
|
||||
// If absent, repetition levels are stored as raw u16 values.
|
||||
CompressiveEncoding rep_compression = 7;
|
||||
// Optional compression algorithm used for the definition buffer.
|
||||
// If absent, definition levels are stored as raw u16 values.
|
||||
CompressiveEncoding def_compression = 8;
|
||||
// Number of values in repetition buffer after decompression.
|
||||
uint64 num_rep_values = 9;
|
||||
// Number of values in definition buffer after decompression.
|
||||
uint64 num_def_values = 10;
|
||||
}
|
||||
|
||||
// A layout where large binary data is encoded externally and only
|
||||
// the descriptions (position + size) are placed in the page
|
||||
//
|
||||
// Repdef information is stored in the descriptions. A description with a size of
|
||||
// 0 and a position of 0 is an empty value. A description with a size of 0 and a
|
||||
// non-zero position is a null value and the position is the repdef value.
|
||||
message BlobLayout {
|
||||
// The inner layout used to store the descriptions
|
||||
PageLayout inner_layout = 1;
|
||||
// The meaning of each repdef layer, used to interpret repdef buffers correctly
|
||||
//
|
||||
// The inner layout's repdef layers will always be 1 all valid item layer
|
||||
repeated RepDefLayer layers = 2;
|
||||
}
|
||||
|
||||
// Describes the structural encoding of a page
|
||||
message PageLayout {
|
||||
oneof layout {
|
||||
// A layout used for pages where the data is small
|
||||
MiniBlockLayout mini_block_layout = 1;
|
||||
// A layout used for pages where all (visible) values are the same scalar value or null.
|
||||
ConstantLayout constant_layout = 2;
|
||||
// A layout used for pages where the data is large
|
||||
FullZipLayout full_zip_layout = 3;
|
||||
// A layout where large binary data is encoded externally
|
||||
// and only the descriptions are put in the page
|
||||
BlobLayout blob_layout = 4;
|
||||
// A sparse structural layout. This variant requires file version 2.3 or later.
|
||||
SparseLayout sparse_layout = 5;
|
||||
}
|
||||
}
|
||||
|
||||
// # Compressive Encodings
|
||||
//
|
||||
// These encodings describe how an array is compressed. An encoding may split an
|
||||
// array into multiple buffers. The buffers can then be compressed further (and split
|
||||
// into yet more buffers). The entire process forms a tree of encodings with the root
|
||||
// of the tree being the initial array and the leaves being the final compressed buffers.
|
||||
//
|
||||
// # Data blocks and buffers
|
||||
//
|
||||
// Data blocks are a simplified version of arrays and represent a collection of buffers grouped
|
||||
// with some kind of interpretation. Data blocks are the input and output of compressive encodings.
|
||||
// There are different kinds of data blocks:
|
||||
// - Fixed width data blocks (e.g. u8, u16, ...)
|
||||
// - Variable width data blocks (e.g. strings, binary)
|
||||
// - Struct data blocks (note: this is for packed structs, normal structs are encoded in the structural encoding)
|
||||
//
|
||||
// In addition, leaf encodings may output "buffers". These are fully compressed buffers of data that
|
||||
// are stored in the page and no longer compressed.
|
||||
|
||||
enum CompressionScheme {
|
||||
COMPRESSION_ALGORITHM_UNSPECIFIED = 0;
|
||||
COMPRESSION_ALGORITHM_LZ4 = 1;
|
||||
COMPRESSION_ALGORITHM_ZSTD = 2;
|
||||
}
|
||||
|
||||
// Compression applied to a single buffer of data
|
||||
//
|
||||
// A buffer is the leaf of the compression tree. Unlike data blocks, which can
|
||||
// be further compressed with a variety of techniques, a buffer cannot be understood
|
||||
// in any particular way.
|
||||
//
|
||||
// A general compression scheme may be applied to a buffer. This is something like
|
||||
// zstd, lz4, etc. The entire buffer is compressed as a single unit. If this happens
|
||||
// then any parent encoding becomes opaque, even if it would normally be transparent.
|
||||
//
|
||||
// This is a leaf, no further compression is applied to the data.
|
||||
message BufferCompression {
|
||||
// A general compression scheme to apply to the buffer
|
||||
CompressionScheme scheme = 1;
|
||||
// The compression level
|
||||
//
|
||||
// Optional, if not present a scheme-specific default value will be used.
|
||||
//
|
||||
// Interpretation of this value depends on the compression scheme. Generally, larger
|
||||
// values indicate more compression at the expense of more CPU time.
|
||||
optional int32 level = 2;
|
||||
}
|
||||
|
||||
// Fixed width items placed contiguously in a single buffer
|
||||
//
|
||||
// This is a leaf encoding, there is no compression applied to the data.
|
||||
//
|
||||
// This is a transparent encoding by definition.
|
||||
//
|
||||
// The input is a fixed-width data block.
|
||||
// The output is a single buffer.
|
||||
message Flat {
|
||||
// the number of bits per value, must be greater than 0, does
|
||||
// not need to be a multiple of 8
|
||||
uint64 bits_per_value = 1;
|
||||
// The compression applied to the data
|
||||
optional BufferCompression data = 2;
|
||||
}
|
||||
|
||||
// Variable width items have the values stored in one buffer and the
|
||||
// offsets are output as a data block that may be further compressed.
|
||||
//
|
||||
// This is a partial leaf encoding. Values are not compressed but
|
||||
// the offsets may be further compressed.
|
||||
//
|
||||
// This is a transparent encoding by definition.
|
||||
//
|
||||
// The input is a variable-width data block.
|
||||
// The output is a single fixed-width data block (the offsets) and
|
||||
// a single buffer (the values)
|
||||
message Variable {
|
||||
// Describes how the offsets data block is compressed
|
||||
CompressiveEncoding offsets = 1;
|
||||
// The compression applied to the values
|
||||
optional BufferCompression values = 2;
|
||||
}
|
||||
|
||||
// Compression algorithm where all values have a constant value (encoded in the description)
|
||||
//
|
||||
// This is a leaf encoding, there is no compression applied to the data.
|
||||
//
|
||||
// The input can be any kind of data block.
|
||||
// There is no output.
|
||||
message Constant {
|
||||
// The value (TODO: define encoding for literals?)
|
||||
optional bytes value = 1;
|
||||
}
|
||||
|
||||
// A compression scheme in which a single fixed-width block is "packed" into
|
||||
// a smaller fixed-width block values where each value has fewer bits.
|
||||
//
|
||||
// This is typically done by throwing away the most significant bits of each value when
|
||||
// those bits are all the same.
|
||||
//
|
||||
// In this scheme the number of bits per value is fixed across the entire buffer and stored
|
||||
// in this message.
|
||||
//
|
||||
// This is a transparent encoding.
|
||||
//
|
||||
// The input is a fixed-width data block.
|
||||
// The output is a single fixed-width data block.
|
||||
message OutOfLineBitpacking {
|
||||
// the number of bits of the uncompressed value. e.g. for a u32, this will be 32
|
||||
uint64 uncompressed_bits_per_value = 1;
|
||||
// The compression used to store the bitpacked values data block
|
||||
CompressiveEncoding values = 3;
|
||||
}
|
||||
|
||||
// Bitpacking variant where the bits per value are stored inline in the chunks themselves
|
||||
//
|
||||
// This variation of bitpacking allows for the number of bits per value to change throughout the
|
||||
// buffer, which makes the compression more robust to outliers.
|
||||
//
|
||||
// This is an opaque encoding.
|
||||
//
|
||||
// The input is a fixed-width data block.
|
||||
// The output is a single buffer.
|
||||
message InlineBitpacking {
|
||||
// the number of bits of the uncompressed value. e.g. for a u32, this will be 32
|
||||
uint64 uncompressed_bits_per_value = 1;
|
||||
// The compression applied to the values
|
||||
optional BufferCompression values = 2;
|
||||
}
|
||||
|
||||
// A compression scheme for variable-width data
|
||||
//
|
||||
// A small dictionary (referred to as a "symbol table") is used to compress the values.
|
||||
// In this scheme there is a single symbol table for the entire page and it is stored in the
|
||||
// encoding description itself.
|
||||
//
|
||||
// This is a transparent encoding.
|
||||
//
|
||||
// The input is a variable-width data block.
|
||||
// The output is a single variable-width data block.
|
||||
message Fsst {
|
||||
// The FSST symbol table
|
||||
bytes symbol_table = 1;
|
||||
// The compression used to store the compressed values data block
|
||||
CompressiveEncoding values = 2;
|
||||
}
|
||||
|
||||
// A compression scheme where common values are stored in a dictionary and the values are
|
||||
// encoded as indices into the dictionary.
|
||||
//
|
||||
// This is an opaque encoding unless the dictionary is considered metadata.
|
||||
//
|
||||
// The input is a any kind of data block.
|
||||
// There are two outputs:
|
||||
// - A data block of the same kind as the input (the dictionary)
|
||||
// - A fixed-width data block containing the indices into the dictionary.
|
||||
message Dictionary {
|
||||
// The compression used to store the indices data block
|
||||
CompressiveEncoding indices = 1;
|
||||
// The compression used to store the dictionary items data block
|
||||
CompressiveEncoding items = 2;
|
||||
// The number of items in the dictionary
|
||||
uint32 num_dictionary_items = 3;
|
||||
}
|
||||
|
||||
// A compression scheme where runs of common values are encoded as a single value and a count
|
||||
//
|
||||
// This is an opaque encoding unless the run lengths are considered metadata.
|
||||
//
|
||||
// The input is a single data block of any kind.
|
||||
// There are two outputs:
|
||||
// - A data block of the same kind as the input (the run values)
|
||||
// - A fixed-width data block containing the lengths of the runs
|
||||
message Rle {
|
||||
// The compression used to store the run values data block
|
||||
CompressiveEncoding values = 1;
|
||||
// The compression used to store the run lengths data block
|
||||
CompressiveEncoding run_lengths = 2;
|
||||
}
|
||||
|
||||
// Converts a fixed-size-list of values into a flattened list of values
|
||||
//
|
||||
// This encoding does not actually compress the data, it just flattens out the FSL layers.
|
||||
//
|
||||
// This is a transparent encoding.
|
||||
//
|
||||
// The input is a single block of fixed-width data (with a wide width and few items)
|
||||
// The output is a single block of fixed-width data (with a narrow width and many items)
|
||||
message FixedSizeList {
|
||||
// The number of items in this layer of FSL
|
||||
uint64 items_per_value = 1;
|
||||
// Whether or not there is a validity buffer
|
||||
bool has_validity = 3;
|
||||
// The compression used to store the flattened values data block
|
||||
CompressiveEncoding values = 2;
|
||||
}
|
||||
|
||||
// Packs a struct containing only fixed-width children into a single fixed-width data block
|
||||
//
|
||||
// The children are concatenated row by row and stored as a single fixed-width buffer. This is
|
||||
// the legacy packed struct representation and remains available for backwards compatibility.
|
||||
message PackedStruct {
|
||||
// The number of bits contributed by each child field in the packed row
|
||||
repeated uint64 bits_per_value = 1;
|
||||
// The compression used to store the packed fixed-width values
|
||||
CompressiveEncoding values = 2;
|
||||
}
|
||||
|
||||
// Variable-width packed struct encoding (2.2 extension)
|
||||
//
|
||||
// Each child value is compressed independently before being transposed into
|
||||
// a row-major layout. This preserves per-field compression boundaries at the
|
||||
// cost of disabling mini-block compression. Readers must prefer this field
|
||||
// when present and fall back to the legacy encoding otherwise.
|
||||
message VariablePackedStruct {
|
||||
// Per-field encoding metadata in struct order
|
||||
repeated FieldEncoding fields = 1;
|
||||
|
||||
// Encoding description for a single child field
|
||||
message FieldEncoding {
|
||||
// Compression applied to individual field values before transposition
|
||||
CompressiveEncoding value = 1;
|
||||
oneof layout {
|
||||
// Bit width of each compressed value (when fixed width)
|
||||
uint64 bits_per_value = 2;
|
||||
// Bit width of the length prefix for variable-width compressed values
|
||||
uint64 bits_per_length = 3;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A compression scheme that wraps the underlying data with general compression
|
||||
//
|
||||
// Note: The application of wrapped compression will depend on the layout of the data.
|
||||
// If we apply it to mini-block data then we compress entire mini-blocks. If we apply
|
||||
// it to full-zip data then we compress each value individually.
|
||||
//
|
||||
// Note: Wrapped compression is somewhat unique at the moment as it is applied to the
|
||||
// output of the inner encoding and not the input like all other compressive encodings.
|
||||
//
|
||||
// Note: General compression can usually be applied in two spots. We can apply
|
||||
// it to individual buffers or we can apply it here, to the entire array.
|
||||
//
|
||||
// For example, let's say we are storing mini-blocks of strings and we are using
|
||||
// FSST and bitpacking the offsets. We have something like this...
|
||||
//
|
||||
// WRAPPED(†3) -> FSST -> VARIABLE -(offsets)-> INLINE_BITPACKING -(data)-> FLAT -> BUFFER (†1)
|
||||
// -(data)-> BUFFER (†2)
|
||||
//
|
||||
// General compression can be applied at †1, †2, or †3 (or any combination of these).
|
||||
//
|
||||
// If we apply it at †1 then we apply it just to the bitpacked offsets
|
||||
// If we apply it at †2 then we apply it just to the FSST compressed data
|
||||
// If we apply it at †3 then we apply it to the entire mini-block (both offsets and data)
|
||||
//
|
||||
// The input is a single data block of any kind.
|
||||
// The output is a single data block of the same kind as the input.
|
||||
message General {
|
||||
// The compression to apply to the values
|
||||
BufferCompression compression = 1;
|
||||
// The compression used to store the output data block
|
||||
CompressiveEncoding values = 3;
|
||||
}
|
||||
|
||||
// A compression scheme where fixed-width values are transposed into a series of byte streams
|
||||
//
|
||||
// This is commonly used for floating point values where the upper bits (the mantissa) have a
|
||||
// significantly different meaning than the lower bits. By splitting the values into byte streams
|
||||
// we group the mantissa bits together and the exponent bits together. The end result is typically
|
||||
// more compressible.
|
||||
//
|
||||
// Note that this encoding is mostly useful when combined with other encodings. It does not do any
|
||||
// compression on its own.
|
||||
//
|
||||
// This is an opaque encoding.
|
||||
//
|
||||
// The input is a fixed-width data block
|
||||
// The output is a single fixed-width data block
|
||||
message ByteStreamSplit {
|
||||
// The compression used to store the values
|
||||
CompressiveEncoding values = 1;
|
||||
}
|
||||
|
||||
// An encoding that compresses a data block into buffers
|
||||
message CompressiveEncoding {
|
||||
oneof compression {
|
||||
Flat flat = 1;
|
||||
Variable variable = 2;
|
||||
Constant constant = 3;
|
||||
OutOfLineBitpacking out_of_line_bitpacking = 4;
|
||||
InlineBitpacking inline_bitpacking = 5;
|
||||
Fsst fsst = 6;
|
||||
Dictionary dictionary = 7;
|
||||
Rle rle = 8;
|
||||
ByteStreamSplit byte_stream_split = 9;
|
||||
General general = 10;
|
||||
FixedSizeList fixed_size_list = 11;
|
||||
PackedStruct packed_struct = 12;
|
||||
VariablePackedStruct variable_packed_struct = 13;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.file;
|
||||
|
||||
// A file descriptor that describes the contents of a Lance file
|
||||
message FileDescriptor {
|
||||
// The schema of the file
|
||||
Schema schema = 1;
|
||||
// The number of rows in the file
|
||||
uint64 length = 2;
|
||||
}
|
||||
|
||||
// A schema which describes the data type of each of the columns
|
||||
message Schema {
|
||||
// All fields in this file, including the nested fields.
|
||||
repeated lance.file.Field fields = 1;
|
||||
// Schema metadata.
|
||||
map<string, bytes> metadata = 5;
|
||||
}
|
||||
|
||||
// Metadata of one Lance file.
|
||||
message Metadata {
|
||||
// 4 was used for StatisticsMetadata in the past, but has been moved to
|
||||
// prevent a bug in older readers.
|
||||
reserved 4;
|
||||
|
||||
// Position of the manifest in the file. If it is zero, the manifest is stored
|
||||
// externally.
|
||||
uint64 manifest_position = 1;
|
||||
|
||||
// Logical offsets of each chunk group, i.e., number of the rows in each
|
||||
// chunk.
|
||||
repeated int32 batch_offsets = 2;
|
||||
|
||||
// The file position that page table is stored.
|
||||
//
|
||||
// A page table is a matrix of N x M x 2, where N = num_fields, and M =
|
||||
// num_batches. Each cell in the table is a pair of <position:int64,
|
||||
// length:int64> of the page. Both position and length are int64 values. The
|
||||
// <position, length> of all the pages in the same column are then
|
||||
// contiguously stored.
|
||||
//
|
||||
// Every field that is a part of the file will have a run in the page table.
|
||||
// This includes struct columns, which will have a run of length 0 since
|
||||
// they don't store any actual data.
|
||||
//
|
||||
// For example, for the column 5 and batch 4, we have:
|
||||
// ```text
|
||||
// position = page_table[5][4][0];
|
||||
// length = page_table[5][4][1];
|
||||
// ```
|
||||
uint64 page_table_position = 3;
|
||||
|
||||
message StatisticsMetadata {
|
||||
// The schema of the statistics.
|
||||
//
|
||||
// This might be empty, meaning there are no statistics. It also might not
|
||||
// contain statistics for every field.
|
||||
repeated Field schema = 1;
|
||||
|
||||
// The field ids of the statistics leaf fields.
|
||||
//
|
||||
// This plays a similar role to the `fields` field in the DataFile message.
|
||||
// Each of these field ids corresponds to a field in the stats_schema. There
|
||||
// is one per column in the stats page table.
|
||||
repeated int32 fields = 2;
|
||||
|
||||
// The file position of the statistics page table
|
||||
//
|
||||
// The page table is a matrix of N x 2, where N = length of stats_fields.
|
||||
// This is the same layout as the main page table, except there is always
|
||||
// only one batch.
|
||||
//
|
||||
// For example, to get the stats column 5, we have:
|
||||
// ```text
|
||||
// position = stats_page_table[5][0];
|
||||
// length = stats_page_table[5][1];
|
||||
// ```
|
||||
uint64 page_table_position = 3;
|
||||
}
|
||||
|
||||
StatisticsMetadata statistics = 5;
|
||||
} // Metadata
|
||||
|
||||
// Supported encodings.
|
||||
enum Encoding {
|
||||
// Invalid encoding.
|
||||
NONE = 0;
|
||||
// Plain encoding.
|
||||
PLAIN = 1;
|
||||
// Var-length binary encoding.
|
||||
VAR_BINARY = 2;
|
||||
// Dictionary encoding.
|
||||
DICTIONARY = 3;
|
||||
// Run-length encoding.
|
||||
RLE = 4;
|
||||
}
|
||||
|
||||
// Dictionary field metadata
|
||||
message Dictionary {
|
||||
/// The file offset for storing the dictionary value.
|
||||
/// It is only valid if encoding is DICTIONARY.
|
||||
///
|
||||
/// The logic type presents the value type of the column, i.e., string value.
|
||||
int64 offset = 1;
|
||||
|
||||
/// The length of dictionary values.
|
||||
int64 length = 2;
|
||||
}
|
||||
|
||||
// Field metadata for a column.
|
||||
message Field {
|
||||
enum Type {
|
||||
PARENT = 0;
|
||||
REPEATED = 1;
|
||||
LEAF = 2;
|
||||
}
|
||||
Type type = 1;
|
||||
|
||||
// Fully qualified name.
|
||||
string name = 2;
|
||||
/// Field Id.
|
||||
///
|
||||
/// See the comment in `DataFile.fields` for how field ids are assigned.
|
||||
int32 id = 3;
|
||||
/// Parent Field ID. If not set, this is a top-level column.
|
||||
int32 parent_id = 4;
|
||||
|
||||
// Logical types, support parameterized Arrow Type.
|
||||
//
|
||||
// PARENT types will always have logical type "struct".
|
||||
//
|
||||
// REPEATED types may have logical types:
|
||||
// * "list"
|
||||
// * "large_list"
|
||||
// * "list.struct"
|
||||
// * "large_list.struct"
|
||||
// The final two are used if the list values are structs, and therefore the
|
||||
// field is both implicitly REPEATED and PARENT.
|
||||
//
|
||||
// LEAF types may have logical types:
|
||||
// * "null"
|
||||
// * "bool"
|
||||
// * "int8" / "uint8"
|
||||
// * "int16" / "uint16"
|
||||
// * "int32" / "uint32"
|
||||
// * "int64" / "uint64"
|
||||
// * "halffloat" / "float" / "double"
|
||||
// * "string" / "large_string"
|
||||
// * "binary" / "large_binary"
|
||||
// * "date32:day"
|
||||
// * "date64:ms"
|
||||
// * "decimal:128:{precision}:{scale}" / "decimal:256:{precision}:{scale}"
|
||||
// * "time:{unit}" / "timestamp:{unit}" / "duration:{unit}", where unit is
|
||||
// "s", "ms", "us", "ns"
|
||||
// * "dict:{value_type}:{index_type}:false"
|
||||
string logical_type = 5;
|
||||
// If this field is nullable.
|
||||
bool nullable = 6;
|
||||
|
||||
// optional field metadata (e.g. extension type name/parameters)
|
||||
map<string, bytes> metadata = 10;
|
||||
|
||||
bool unenforced_primary_key = 12;
|
||||
|
||||
// Position of this field in the primary key (1-based).
|
||||
// 0 means the field is part of the primary key but uses schema field id for ordering.
|
||||
// When set to a positive value, primary key fields are ordered by this position.
|
||||
uint32 unenforced_primary_key_position = 13;
|
||||
|
||||
// Reserved for future use. Use unenforced_clustering_key_position instead.
|
||||
bool unenforced_clustering_key = 14;
|
||||
|
||||
// Position of this field in the clustering key (1-based).
|
||||
// 0 means the field is not part of the clustering key.
|
||||
uint32 unenforced_clustering_key_position = 15;
|
||||
|
||||
// DEPRECATED ----------------------------------------------------------------
|
||||
|
||||
// Deprecated: Only used in V1 file format. V2 uses variable encodings defined
|
||||
// per page.
|
||||
//
|
||||
// The global encoding to use for this field.
|
||||
Encoding encoding = 7;
|
||||
|
||||
// Deprecated: Only used in V1 file format. V2 dynamically chooses when to
|
||||
// do dictionary encoding and keeps the dictionary in the data files.
|
||||
//
|
||||
// The file offset for storing the dictionary value.
|
||||
// It is only valid if encoding is DICTIONARY.
|
||||
//
|
||||
// The logic type presents the value type of the column, i.e., string value.
|
||||
Dictionary dictionary = 8;
|
||||
|
||||
// Deprecated: optional extension type name, use metadata field
|
||||
// ARROW:extension:name
|
||||
string extension_name = 9;
|
||||
|
||||
// Field number 11 was previously `string storage_class`.
|
||||
// Keep it reserved so older manifests remain compatible while new writers
|
||||
// avoid reusing the slot.
|
||||
reserved 11;
|
||||
reserved "storage_class";
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.file.v2;
|
||||
|
||||
import "google/protobuf/any.proto";
|
||||
import "google/protobuf/empty.proto";
|
||||
|
||||
// # Lance v2.X File Format
|
||||
//
|
||||
// The Lance file format is a barebones format for serializing columnar data
|
||||
// into a file.
|
||||
//
|
||||
// * Each Lance file contains between 0 and 4Gi columns
|
||||
// * Each column contains between 0 and 4Gi pages
|
||||
// * Each page contains between 0 and 2^64 items
|
||||
// * Different pages within a column can have different items counts
|
||||
// * Columns may have up to 2^64 items
|
||||
// * Different columns within a file can have different item counts
|
||||
//
|
||||
// The Lance file format does not have any notion of a type system or schemas.
|
||||
// From the perspective of the file format all data is arbitrary buffers of
|
||||
// bytes with an extensible metadata block to describe the data. It is up to
|
||||
// the user to interpret these bytes meaningfully.
|
||||
//
|
||||
// Data buffers are written to the file first. These data buffers can be
|
||||
// referenced from three different places in the file:
|
||||
//
|
||||
// * Page encodings can reference data buffers. This is the most common way
|
||||
// that actual data is stored.
|
||||
// * Column encodings can reference data buffers. For example, a column encoding
|
||||
// may reference data buffer(s) containing statistics or dictionaries.
|
||||
// * Finally, the global buffer offset table can reference data buffers. This
|
||||
// is useful for storing data that is shared across multiple columns.
|
||||
// This is also useful for global file metadata (e.g. a schema that describes
|
||||
// the file)
|
||||
//
|
||||
// ## File Layout
|
||||
//
|
||||
// Note: the number of buffers (BN) is independent of the number of columns (CN)
|
||||
// and pages.
|
||||
//
|
||||
// Buffers often need to be aligned. 64-byte alignment is common when
|
||||
// working with SIMD operations. 4096-byte alignment is common when
|
||||
// working with direct I/O. In order to ensure these buffers are aligned
|
||||
// writers may need to insert padding before the buffers.
|
||||
//
|
||||
// If direct I/O is required then most (but not all) fields described
|
||||
// below must be sector aligned. We have marked these fields with an
|
||||
// asterisk for clarity. Readers should assume there will be optional
|
||||
// padding inserted before these fields.
|
||||
//
|
||||
// All footer fields are unsigned integers written with little endian
|
||||
// byte order.
|
||||
//
|
||||
// ├──────────────────────────────────┤
|
||||
// | Data Pages |
|
||||
// | Data Buffer 0* |
|
||||
// | ... |
|
||||
// | Data Buffer BN* |
|
||||
// ├──────────────────────────────────┤
|
||||
// | Column Metadatas |
|
||||
// | |A| Column 0 Metadata* |
|
||||
// | Column 1 Metadata* |
|
||||
// | ... |
|
||||
// | Column CN Metadata* |
|
||||
// ├──────────────────────────────────┤
|
||||
// | Column Metadata Offset Table |
|
||||
// | |B| Column 0 Metadata Position* |
|
||||
// | Column 0 Metadata Size |
|
||||
// | ... |
|
||||
// | Column CN Metadata Position |
|
||||
// | Column CN Metadata Size |
|
||||
// ├──────────────────────────────────┤
|
||||
// | Global Buffers Offset Table |
|
||||
// | |C| Global Buffer 0 Position* |
|
||||
// | Global Buffer 0 Size |
|
||||
// | ... |
|
||||
// | Global Buffer GN Position |
|
||||
// | Global Buffer GN Size |
|
||||
// ├──────────────────────────────────┤
|
||||
// | Footer |
|
||||
// | A u64: Offset to column meta 0 |
|
||||
// | B u64: Offset to CMO table |
|
||||
// | C u64: Offset to GBO table |
|
||||
// | u32: Number of global bufs |
|
||||
// | u32: Number of columns |
|
||||
// | u16: Major version |
|
||||
// | u16: Minor version |
|
||||
// | "LANC" |
|
||||
// ├──────────────────────────────────┤
|
||||
//
|
||||
// File Layout-End
|
||||
//
|
||||
// ## Data Pages
|
||||
//
|
||||
// A lot of flexibility is provided in how data is stored. A page's buffers do
|
||||
// not strictly need to be contiguous on the disk. However, it is recommended
|
||||
// that buffers within a page be grouped together for best performance.
|
||||
//
|
||||
// Data pages should be large. The only time a page should be written to disk
|
||||
// is when the writer needs to flush the page to disk because it has accumulated
|
||||
// too much data. Pages are not read in sequential order and if pages are too
|
||||
// small then the seek overhead (or request overhead) will be problematic. We
|
||||
// generally advise that pages be at least 8MB or larger.
|
||||
//
|
||||
// ## Encodings
|
||||
//
|
||||
// Specific encodings are not part of this minimal format. They are provided
|
||||
// by extensions. Readers and writers should be designed so that encodings can
|
||||
// be easily added and removed. Ideally, they should allow for this without
|
||||
// requiring recompilation through some kind of plugin system.
|
||||
|
||||
// The deferred encoding is used to place the encoding itself in a different
|
||||
// part of the file. This is most commonly used to allow encodings to be shared
|
||||
// across different columns. For example, when writing a file with thousands of
|
||||
// columns, where many pages have the exact same encoding, it can be useful
|
||||
// to cut down on the size of the metadata by using a deferred encoding.
|
||||
message DeferredEncoding {
|
||||
// Location of the buffer containing the encoding.
|
||||
//
|
||||
// * If sharing encodings across columns then this will be in a global buffer
|
||||
// * If sharing encodings across pages within a column this could be in a
|
||||
// column metadata buffer.
|
||||
// * This could also be a page buffer if the encoding is not shared, needs
|
||||
// to be written before the file ends, and the encoding is too large to load
|
||||
// unless we first determine the page needs to be read. This combination
|
||||
// seems unusual.
|
||||
uint64 buffer_location = 1;
|
||||
uint64 buffer_length = 2;
|
||||
}
|
||||
|
||||
// The encoding is placed directly in the metadata section
|
||||
message DirectEncoding {
|
||||
// The bytes that make up the encoding embedded directly in the metadata
|
||||
//
|
||||
// This is the most common approach.
|
||||
bytes encoding = 1;
|
||||
}
|
||||
|
||||
// An encoding stores the information needed to decode a column or page
|
||||
//
|
||||
// For example, it could describe if the page is using bit packing, and how many bits
|
||||
// there are in each individual value.
|
||||
//
|
||||
// At the column level it can be used to wrap columns with dictionaries or statistics.
|
||||
message Encoding {
|
||||
oneof location {
|
||||
// The encoding is stored elsewhere and not part of this protobuf message
|
||||
DeferredEncoding indirect = 1;
|
||||
// The encoding is stored within this protobuf message
|
||||
DirectEncoding direct = 2;
|
||||
// There is no encoding information
|
||||
google.protobuf.Empty none = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// ## Metadata
|
||||
|
||||
// Each column has a metadata block that is placed at the end of the file.
|
||||
// These may be read individually to allow for column projection.
|
||||
message ColumnMetadata {
|
||||
|
||||
// This describes a page of column data.
|
||||
message Page {
|
||||
// The file offsets for each of the page buffers
|
||||
//
|
||||
// The number of buffers is variable and depends on the encoding. There
|
||||
// may be zero buffers (e.g. constant encoded data) in which case this
|
||||
// could be empty.
|
||||
repeated uint64 buffer_offsets = 1;
|
||||
// The size (in bytes) of each of the page buffers
|
||||
//
|
||||
// This field will have the same length as `buffer_offsets` and
|
||||
// may be empty.
|
||||
repeated uint64 buffer_sizes = 2;
|
||||
// Logical length (e.g. # rows) of the page
|
||||
uint64 length = 3;
|
||||
// The encoding used to encode the page
|
||||
Encoding encoding = 4;
|
||||
// The priority of the page
|
||||
//
|
||||
// For tabular data this will be the top-level row number of the first row
|
||||
// in the page (and top-level rows should not split across pages).
|
||||
uint64 priority = 5;
|
||||
}
|
||||
// Encoding information about the column itself. This typically describes
|
||||
// how to interpret the column metadata buffers. For example, it could
|
||||
// describe how statistics or dictionaries are stored in the column metadata.
|
||||
Encoding encoding = 1;
|
||||
// The pages in the column
|
||||
repeated Page pages = 2;
|
||||
// The file offsets of each of the column metadata buffers
|
||||
//
|
||||
// There may be zero buffers.
|
||||
repeated uint64 buffer_offsets = 3;
|
||||
// The size (in bytes) of each of the column metadata buffers
|
||||
//
|
||||
// This field will have the same length as `buffer_offsets` and
|
||||
// may be empty.
|
||||
repeated uint64 buffer_sizes = 4;
|
||||
} // Metadata-End
|
||||
|
||||
// ## Where is the rest?
|
||||
//
|
||||
// This file format is extremely minimal. It is a building block for
|
||||
// creating more useful readers and writers and not terribly useful by itself.
|
||||
// Other protobuf files will describe how this can be extended.
|
||||
@@ -0,0 +1,99 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.datafusion;
|
||||
|
||||
import "table_identifier.proto";
|
||||
|
||||
message U64Range {
|
||||
uint64 start = 1;
|
||||
uint64 end = 2;
|
||||
}
|
||||
|
||||
message ProjectionProto {
|
||||
repeated int32 field_ids = 1;
|
||||
bool with_row_id = 2;
|
||||
bool with_row_addr = 3;
|
||||
bool with_row_last_updated_at_version = 4;
|
||||
bool with_row_created_at_version = 5;
|
||||
BlobHandlingProto blob_handling = 6;
|
||||
}
|
||||
|
||||
message BlobHandlingProto {
|
||||
oneof mode {
|
||||
// All blobs read as binary
|
||||
bool all_binary = 1;
|
||||
// Blobs as descriptions, other binary as binary (default)
|
||||
bool blobs_descriptions = 2;
|
||||
// All binary columns as descriptions
|
||||
bool all_descriptions = 3;
|
||||
// Specific blobs read as binary, rest as descriptions (non-blob binary stays binary)
|
||||
FieldIdSet some_blobs_binary = 4;
|
||||
// Specific columns as binary, all other binary as descriptions
|
||||
FieldIdSet some_binary = 5;
|
||||
}
|
||||
}
|
||||
|
||||
message FieldIdSet {
|
||||
repeated uint32 field_ids = 1;
|
||||
}
|
||||
|
||||
message FilteredReadThreadingModeProto {
|
||||
oneof mode {
|
||||
uint64 one_partition_multiple_threads = 1;
|
||||
uint64 multiple_partitions = 2;
|
||||
}
|
||||
}
|
||||
|
||||
// Serializable form of FilteredReadOptions.
|
||||
message FilteredReadOptionsProto {
|
||||
optional U64Range scan_range_before_filter = 1;
|
||||
optional U64Range scan_range_after_filter = 2;
|
||||
bool with_deleted_rows = 3;
|
||||
optional uint32 batch_size = 4;
|
||||
optional uint64 fragment_readahead = 5;
|
||||
repeated uint64 fragment_ids = 6;
|
||||
ProjectionProto projection = 7;
|
||||
optional bytes refine_filter_substrait = 8;
|
||||
optional bytes full_filter_substrait = 9;
|
||||
FilteredReadThreadingModeProto threading_mode = 10;
|
||||
optional uint64 io_buffer_size_bytes = 11;
|
||||
// Arrow IPC schema for decoding Substrait filters (may be wider than projection).
|
||||
optional bytes filter_schema_ipc = 12;
|
||||
}
|
||||
|
||||
// Serializable form of FilteredReadPlan (planned/distributed mode).
|
||||
// RowAddrTreeMap serialized via its built-in serialize_into/deserialize_from.
|
||||
// Per-fragment filters are Substrait-encoded and deduplicated.
|
||||
message FilteredReadPlanProto {
|
||||
bytes row_addr_tree_map = 1;
|
||||
optional U64Range scan_range_after_filter = 2;
|
||||
// Arrow IPC schema for decoding Substrait filters (matches the schema used at encode time).
|
||||
optional bytes filter_schema_ipc = 3;
|
||||
// Per-fragment filter mapping. Key is fragment id, value is a list index into
|
||||
// filter_expressions. Multiple fragments can share the same list index when
|
||||
// they have the same filter, avoiding duplicate Substrait encoding.
|
||||
map<uint32, uint32> fragment_filter_ids = 4;
|
||||
// Deduplicated Substrait-encoded filter expressions. Each entry is referenced
|
||||
// by one or more values in fragment_filter_ids.
|
||||
repeated bytes filter_expressions = 5;
|
||||
}
|
||||
|
||||
// Top-level wrapper for FilteredReadExec serialization.
|
||||
message FilteredReadExecProto {
|
||||
TableIdentifier table = 1;
|
||||
FilteredReadOptionsProto options = 2;
|
||||
// FilteredRead has two modes
|
||||
// Plan-then-execute (distributed): The planner creates a FilteredReadPlan and sends it to a remote executor.
|
||||
// Plan-and-execute (local): The executor creates the plan itself at execution time.
|
||||
optional FilteredReadPlanProto plan = 3;
|
||||
// Note: FilteredReadExec.index_input (child ExecutionPlan) is NOT serialized here.
|
||||
// DataFusion's PhysicalExtensionCodec handles child plans automatically: it walks
|
||||
// the plan tree via children() / with_new_children(), serializes each node, and
|
||||
// passes deserialized children back as the `inputs` parameter in try_decode.
|
||||
// This means any ExecutionPlan in the tree (including index_input) must also
|
||||
// implement try_encode/try_decode in the PhysicalExtensionCodec.
|
||||
// TODO: implement serialize/deserialize for lance-specific index input ExecutionPlans.
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.index.pb;
|
||||
|
||||
import "google/protobuf/any.proto";
|
||||
|
||||
// The type of an index.
|
||||
enum IndexType {
|
||||
// Vector index
|
||||
VECTOR = 0;
|
||||
}
|
||||
|
||||
message Index {
|
||||
// The unique index name in the dataset.
|
||||
string name = 1;
|
||||
|
||||
// Columns to be used to build the index.
|
||||
repeated string columns = 2;
|
||||
|
||||
// The version of the dataset this index was built from.
|
||||
uint64 dataset_version = 3;
|
||||
|
||||
// The [`IndexType`] of the index.
|
||||
IndexType index_type = 4;
|
||||
|
||||
/// Index implementation details.
|
||||
oneof implementation {
|
||||
VectorIndex vector_index = 5;
|
||||
}
|
||||
}
|
||||
|
||||
message Tensor {
|
||||
enum DataType {
|
||||
BFLOAT16 = 0;
|
||||
FLOAT16 = 1;
|
||||
FLOAT32 = 2;
|
||||
FLOAT64 = 3;
|
||||
UINT8 = 4;
|
||||
UINT16 = 5;
|
||||
UINT32 = 6;
|
||||
UINT64 = 7;
|
||||
}
|
||||
|
||||
DataType data_type = 1;
|
||||
|
||||
// Data shape, [dim1, dim2, ...]
|
||||
repeated uint32 shape = 2;
|
||||
|
||||
// Data buffer
|
||||
bytes data = 3;
|
||||
}
|
||||
|
||||
// Inverted Index File Metadata.
|
||||
message IVF {
|
||||
// Centroids of partitions. `dimension * num_partitions` of float32s.
|
||||
//
|
||||
// Deprecated, use centroids_tensor instead.
|
||||
repeated float centroids = 1; // [deprecated = true];
|
||||
|
||||
// File offset of each partition.
|
||||
repeated uint64 offsets = 2;
|
||||
|
||||
// Number of records in the partition.
|
||||
repeated uint32 lengths = 3;
|
||||
|
||||
// Tensor of centroids. `num_partitions * dimension` of float32s.
|
||||
Tensor centroids_tensor = 4;
|
||||
|
||||
// KMeans loss.
|
||||
optional double loss = 5;
|
||||
}
|
||||
|
||||
// Product Quantization.
|
||||
message PQ {
|
||||
// The number of bits to present a centroid.
|
||||
uint32 num_bits = 1;
|
||||
|
||||
// Number of sub vectors.
|
||||
uint32 num_sub_vectors = 2;
|
||||
|
||||
// Vector dimension
|
||||
uint32 dimension = 3;
|
||||
|
||||
// Codebook. `dimension * 2 ^ num_bits` of float32s.
|
||||
repeated float codebook = 4;
|
||||
|
||||
// Tensor of codebook. `2 ^ num_bits * dimension` of floats.
|
||||
Tensor codebook_tensor = 5;
|
||||
}
|
||||
|
||||
// Transform type
|
||||
enum TransformType {
|
||||
OPQ = 0;
|
||||
}
|
||||
|
||||
// A transform matrix to apply to a vector or vectors.
|
||||
message Transform {
|
||||
// The file offset the matrix is stored
|
||||
uint64 position = 1;
|
||||
|
||||
// Data shape of the matrix, [rows, cols].
|
||||
repeated uint32 shape = 2;
|
||||
|
||||
// Transform type.
|
||||
TransformType type = 3;
|
||||
}
|
||||
|
||||
// Flat Index
|
||||
message Flat {}
|
||||
|
||||
// DiskAnn Index
|
||||
message DiskAnn {
|
||||
// Graph spec version
|
||||
uint32 spec = 1;
|
||||
|
||||
// Graph file
|
||||
string filename = 2;
|
||||
|
||||
// r parameter
|
||||
uint32 r = 3;
|
||||
|
||||
// alpha parameter
|
||||
float alpha = 4;
|
||||
|
||||
// L parameter
|
||||
uint32 L = 5;
|
||||
|
||||
/// Entry points to the graph
|
||||
repeated uint64 entries = 6;
|
||||
}
|
||||
|
||||
// One stage in the vector index pipeline.
|
||||
message VectorIndexStage {
|
||||
oneof stage {
|
||||
// Flat index
|
||||
Flat flat = 1;
|
||||
// `IVF` - Inverted File
|
||||
IVF ivf = 2;
|
||||
// Product Quantization
|
||||
PQ pq = 3;
|
||||
// Transformer
|
||||
Transform transform = 4;
|
||||
// DiskANN
|
||||
DiskAnn diskann = 5;
|
||||
}
|
||||
}
|
||||
|
||||
// Metric Type for Vector Index
|
||||
enum VectorMetricType {
|
||||
// L2 (Euclidean) Distance
|
||||
L2 = 0;
|
||||
|
||||
// Cosine Distance
|
||||
Cosine = 1;
|
||||
|
||||
// Dot Product
|
||||
Dot = 2;
|
||||
|
||||
// Hamming Distance
|
||||
Hamming = 3;
|
||||
}
|
||||
|
||||
// Vector Index Metadata
|
||||
message VectorIndex {
|
||||
// Index specification version.
|
||||
uint32 spec_version = 1;
|
||||
|
||||
// Vector dimension;
|
||||
uint32 dimension = 2;
|
||||
|
||||
// Composed vector index stages.
|
||||
//
|
||||
// For example, `IVF_PQ` index type can be expressed as:
|
||||
//
|
||||
// ```text
|
||||
// let stages = vec![Ivf{}, PQ{num_bits: 8, num_sub_vectors: 16}]
|
||||
// ```
|
||||
repeated VectorIndexStage stages = 3;
|
||||
|
||||
// Vector distance metrics type
|
||||
VectorMetricType metric_type = 4;
|
||||
}
|
||||
|
||||
// Details for vector indexes, stored in the manifest's index_details field.
|
||||
message VectorIndexDetails {
|
||||
VectorMetricType metric_type = 1;
|
||||
|
||||
// The target number of vectors per partition.
|
||||
// 0 means unset.
|
||||
uint64 target_partition_size = 2;
|
||||
|
||||
// Optional HNSW index configuration. If set, the index has an HNSW layer.
|
||||
optional HnswParameters hnsw_index_config = 3;
|
||||
|
||||
message ProductQuantization {
|
||||
uint32 num_bits = 1;
|
||||
uint32 num_sub_vectors = 2;
|
||||
}
|
||||
message ScalarQuantization {
|
||||
uint32 num_bits = 1;
|
||||
}
|
||||
message RabitQuantization {
|
||||
enum RotationType {
|
||||
FAST = 0;
|
||||
MATRIX = 1;
|
||||
}
|
||||
uint32 num_bits = 1;
|
||||
RotationType rotation_type = 2;
|
||||
}
|
||||
|
||||
// No quantization; vectors are stored as-is.
|
||||
message FlatCompression {}
|
||||
|
||||
oneof compression {
|
||||
ProductQuantization pq = 4;
|
||||
ScalarQuantization sq = 5;
|
||||
RabitQuantization rq = 6;
|
||||
FlatCompression flat = 8;
|
||||
}
|
||||
|
||||
// Runtime hints: optional build preferences that don't affect index structure.
|
||||
// Keys use reverse-DNS namespacing (e.g., "lance.ivf.max_iters", "lancedb.accelerator").
|
||||
// Unrecognized keys must be silently ignored by all runtimes.
|
||||
map<string, string> runtime_hints = 9;
|
||||
}
|
||||
|
||||
// Hierarchical Navigable Small World (HNSW) parameters, used as an optional configuration for IVF indexes.
|
||||
message HnswParameters {
|
||||
// The maximum number of outgoing edges per node in the HNSW graph. Higher values
|
||||
// means more connections, better recall, but more memory and slower builds.
|
||||
// Referred to as "M" in the HNSW literature.
|
||||
uint32 max_connections = 1;
|
||||
// "construction exploration factor": The size of the dynamic list used during
|
||||
// index construction.
|
||||
uint32 construction_ef = 2;
|
||||
// The maximum number of levels in the HNSW graph.
|
||||
uint32 max_level = 3;
|
||||
}
|
||||
|
||||
message JsonIndexDetails {
|
||||
string path = 1;
|
||||
google.protobuf.Any target_details = 2;
|
||||
}
|
||||
message BloomFilterIndexDetails {}
|
||||
|
||||
message RTreeIndexDetails {}
|
||||
|
||||
message FMIndexDetails {}
|
||||
@@ -0,0 +1,104 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.table;
|
||||
|
||||
// NOTE: Do *NOT* add new index details here. Add them to the index.proto file instead.
|
||||
// This file is in the lance.table package namespace while the index.proto file is in the
|
||||
// lance.index package namespace.
|
||||
//
|
||||
// These are only here for forward compatibility. Older versions of Lance expect btree indexes
|
||||
// to have lance.table in the package namespace.
|
||||
//
|
||||
// If you need to modify these messages (e.g. to add new fields to btree or bitmap) then
|
||||
// it is ok to modify them here.
|
||||
|
||||
// Currently many of these are empty messages because all needed details are either hard-coded (e.g.
|
||||
// filenames) or stored in the index itself. However, we may want to add more details in the
|
||||
// future, in particular we can add details that may be useful for planning queries (e.g. don't
|
||||
// force us to load the index until we know we can make use of it)
|
||||
|
||||
message BTreeIndexDetails {}
|
||||
message BitmapIndexDetails {}
|
||||
message LabelListIndexDetails {}
|
||||
message NGramIndexDetails {}
|
||||
message ZoneMapIndexDetails {
|
||||
// Number of rows per zone. Optional for backwards compatibility: absent on
|
||||
// datasets written before this field was added. When absent, no seed writer
|
||||
// is created for the index.
|
||||
optional uint64 rows_per_zone = 1;
|
||||
// Whether seed-based incremental updates are enabled for this index.
|
||||
// On-disk semantics: absent means seeds are disabled (old datasets written
|
||||
// before this field was added). Present false means explicitly disabled.
|
||||
// Present true means seeds are enabled: the index will embed per-fragment
|
||||
// seed buffers in data files and harvest them during incremental updates
|
||||
// to skip full column scans.
|
||||
// Creation-time default: index creation code sets this to true for
|
||||
// variable-length types (strings, binary) and fixed-width types wider than
|
||||
// 8 bytes, and to false for narrow fixed-width types (e.g. Int64, Float64).
|
||||
optional bool use_seeds = 2;
|
||||
// Whether this index tracks exact null row addresses in a separate bitmap.
|
||||
// Absent or false means legacy format: null positions are not tracked and
|
||||
// IS NULL searches fall back to approximate zone-level statistics. Present
|
||||
// true means IS NULL is exact and IS NOT NULL can be answered without a
|
||||
// full scan.
|
||||
optional bool has_null_bitmap = 3;
|
||||
}
|
||||
message InvertedIndexDetails {
|
||||
enum DocumentGranularity {
|
||||
ROW = 0;
|
||||
LIST_ELEMENT = 1;
|
||||
}
|
||||
|
||||
message CodeTokenizerConfig {
|
||||
// Split one lexical identifier into subwords, e.g. getUserName ->
|
||||
// get/user/name.
|
||||
bool split_identifiers = 1;
|
||||
// Split identifier subwords across letter/number boundaries, e.g.
|
||||
// HTML2JSON -> html/2/json. An absent value uses the code tokenizer default;
|
||||
// a present value records the explicit index-time choice.
|
||||
optional bool split_on_numerics = 2;
|
||||
// Keep the complete lexical identifier in addition to subwords, e.g.
|
||||
// user_name plus user/name. An absent value uses the code tokenizer default;
|
||||
// a present value records the explicit index-time choice.
|
||||
optional bool preserve_original = 3;
|
||||
// Index operator tokens such as "::", "->", and "!=". Operators are not
|
||||
// indexed by default because they are often high-frequency noise.
|
||||
bool index_operators = 4;
|
||||
}
|
||||
|
||||
// Lexical tokenizer used after document-level text extraction. This is an
|
||||
// implementation component such as "simple", "icu", "ngram", or "code".
|
||||
// Input-time analyzer profiles are expanded into this field and the concrete
|
||||
// options below before these details are persisted.
|
||||
// Marking this field as optional as old versions of the index store blank details and we
|
||||
// need to make sure we have a proper optional field to detect this.
|
||||
optional string base_tokenizer = 1;
|
||||
string language = 2;
|
||||
bool with_position = 3;
|
||||
optional uint32 max_token_length = 4;
|
||||
bool lower_case = 5;
|
||||
bool stem = 6;
|
||||
bool remove_stop_words = 7;
|
||||
bool ascii_folding = 8;
|
||||
uint32 min_ngram_length = 9;
|
||||
uint32 max_ngram_length = 10;
|
||||
bool prefix_only = 11;
|
||||
// Number of documents per compressed posting block. An absent value means
|
||||
// the index predates this field and must use the legacy block size of 128.
|
||||
// A present value records the block size used by the index; 256 is valid
|
||||
// with format versions 3 and 4.
|
||||
optional uint32 block_size = 12;
|
||||
// Options for base_tokenizer = "code". Presence records the code tokenizer
|
||||
// configuration used to build the index; absence means there is no
|
||||
// code-specific configuration to apply.
|
||||
CodeTokenizerConfig code_config = 13;
|
||||
// The logical FTS document boundary. The protobuf default preserves the
|
||||
// legacy row-document behavior when this field is absent.
|
||||
DocumentGranularity document_granularity = 14;
|
||||
// The posting-list payload format. This is separate from index_version,
|
||||
// which identifies the overall inverted-index layout.
|
||||
optional uint32 posting_format_version = 15;
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
@@ -0,0 +1,113 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.table;
|
||||
// TODO: what would it take to store this in a LanceV2 file?
|
||||
// Or would flatbuffers be better for this?
|
||||
|
||||
/// A sequence of row IDs. This is split up into one or more segments,
|
||||
/// each of which can be encoded in different ways. The encodings are optimized
|
||||
/// for values that are sorted, which will often be the case with row ids.
|
||||
/// They also have optimized forms depending on how sparse the values are.
|
||||
message RowIdSequence {
|
||||
repeated U64Segment segments = 1;
|
||||
}
|
||||
|
||||
/// Different ways to encode a sequence of u64 values.
|
||||
message U64Segment {
|
||||
/// A range of u64 values.
|
||||
message Range {
|
||||
/// The start of the range, inclusive.
|
||||
uint64 start = 1;
|
||||
/// The end of the range, exclusive.
|
||||
uint64 end = 2;
|
||||
}
|
||||
|
||||
/// A range of u64 values with holes.
|
||||
message RangeWithHoles {
|
||||
/// The start of the range, inclusive.
|
||||
uint64 start = 1;
|
||||
/// The end of the range, exclusive.
|
||||
uint64 end = 2;
|
||||
/// The holes in the range, as a sorted array of values;
|
||||
/// Binary search can be used to check whether a value is a hole and should
|
||||
/// be skipped. This can also be used to count the number of holes before a
|
||||
/// given value, if you need to find the logical offset of a value in the
|
||||
/// segment.
|
||||
EncodedU64Array holes = 3;
|
||||
}
|
||||
|
||||
/// A range of u64 values with a bitmap.
|
||||
message RangeWithBitmap {
|
||||
/// The start of the range, inclusive.
|
||||
uint64 start = 1;
|
||||
/// The end of the range, exclusive.
|
||||
uint64 end = 2;
|
||||
/// A bitmap of the values in the range. The bitmap is a sequence of bytes,
|
||||
/// where each byte represents 8 values. The first byte represents values
|
||||
/// start to start + 7, the second byte represents values start + 8 to
|
||||
/// start + 15, and so on. The most significant bit of each byte represents
|
||||
/// the first value in the range, and the least significant bit represents
|
||||
/// the last value in the range. If the bit is set, the value is in the
|
||||
/// range; if it is not set, the value is not in the range.
|
||||
bytes bitmap = 3;
|
||||
}
|
||||
|
||||
oneof segment {
|
||||
/// When the values are sorted and contiguous.
|
||||
Range range = 1;
|
||||
/// When the values are sorted but have a few gaps.
|
||||
RangeWithHoles range_with_holes = 2;
|
||||
/// When the values are sorted but have many gaps.
|
||||
RangeWithBitmap range_with_bitmap = 3;
|
||||
/// When the values are sorted but are sparse.
|
||||
EncodedU64Array sorted_array = 4;
|
||||
/// A general array of values, which is not sorted.
|
||||
EncodedU64Array array = 5;
|
||||
}
|
||||
} // RowIdSegment
|
||||
|
||||
/// A basic bitpacked array of u64 values.
|
||||
message EncodedU64Array {
|
||||
message U16Array {
|
||||
uint64 base = 1;
|
||||
/// The deltas are stored as 16-bit unsigned integers.
|
||||
/// (protobuf doesn't support 16-bit integers, so we use bytes instead)
|
||||
bytes offsets = 2;
|
||||
}
|
||||
|
||||
message U32Array {
|
||||
uint64 base = 1;
|
||||
/// The deltas are stored as 32-bit unsigned integers.
|
||||
/// (we use bytes instead of uint32 to avoid overhead of varint encoding)
|
||||
bytes offsets = 2;
|
||||
}
|
||||
|
||||
message U64Array {
|
||||
/// (We use bytes instead of uint64 to avoid overhead of varint encoding)
|
||||
bytes values = 2;
|
||||
}
|
||||
|
||||
oneof array {
|
||||
U16Array u16_array = 1;
|
||||
U32Array u32_array = 2;
|
||||
U64Array u64_array = 3;
|
||||
}
|
||||
}
|
||||
|
||||
/// A sequence of dataset versions. Similar to RowIdSequence but tracks
|
||||
/// version runs. It uses RLE (Run-Length Encoding) to efficiently
|
||||
// represent consecutive rows with the same version.
|
||||
message RowDatasetVersionSequence {
|
||||
repeated RowDatasetVersionRun runs = 1;
|
||||
}
|
||||
|
||||
/// A run of rows with the same version.
|
||||
message RowDatasetVersionRun {
|
||||
/// The number of consecutive rows with the same version.
|
||||
U64Segment span = 1;
|
||||
|
||||
uint64 version = 2;
|
||||
}
|
||||
@@ -0,0 +1,805 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.table;
|
||||
|
||||
import "google/protobuf/any.proto";
|
||||
import "google/protobuf/timestamp.proto";
|
||||
import "file.proto";
|
||||
|
||||
/*
|
||||
|
||||
Format:
|
||||
|
||||
+----------------------------------------+
|
||||
| Encoded Column 0, Chunk 0 |
|
||||
...
|
||||
| Encoded Column M, Chunk N - 1 |
|
||||
| Encoded Column M, Chunk N |
|
||||
| Indices ... |
|
||||
| Chunk Position (M x N x 8) |
|
||||
| Manifest (Optional) |
|
||||
| Metadata |
|
||||
| i64: metadata position |
|
||||
| MAJOR_VERSION | MINOR_VERSION | "LANC" |
|
||||
+----------------------------------------+
|
||||
*/
|
||||
|
||||
// UUID type. encoded as 16 bytes.
|
||||
message UUID {
|
||||
bytes uuid = 1;
|
||||
}
|
||||
|
||||
// Manifest is a global section shared between all the files.
|
||||
message Manifest {
|
||||
// All fields of the dataset, including the nested fields.
|
||||
repeated lance.file.Field fields = 1;
|
||||
|
||||
// Schema metadata.
|
||||
map<string, bytes> schema_metadata = 5;
|
||||
|
||||
// Fragments of the dataset.
|
||||
repeated DataFragment fragments = 2;
|
||||
|
||||
// Snapshot version number.
|
||||
uint64 version = 3;
|
||||
|
||||
// The file position of the version auxiliary data.
|
||||
// * It is not inheritable between versions.
|
||||
// * It is not loaded by default during query.
|
||||
uint64 version_aux_data = 4;
|
||||
|
||||
message WriterVersion {
|
||||
// The name of the library that created this file.
|
||||
string library = 1;
|
||||
// The version of the library that created this file. Because we cannot assume
|
||||
// that the library is semantically versioned, this is a string. However, if it
|
||||
// is semantically versioned, it should be a valid semver string without any 'v'
|
||||
// prefix. For example: `2.0.0`, `2.0.0-rc.1`.
|
||||
//
|
||||
// For forward compatibility with older readers, when writing new manifests this
|
||||
// field should contain only the core version (major.minor.patch) without any
|
||||
// prerelease or build metadata. The prerelease/build info should be stored in
|
||||
// the separate prerelease and build_metadata fields instead.
|
||||
string version = 2;
|
||||
// Optional semver prerelease identifier.
|
||||
//
|
||||
// This field stores the prerelease portion of a semantic version separately
|
||||
// from the core version number. For example, if the full version is "2.0.0-rc.1",
|
||||
// the version field would contain "2.0.0" and prerelease would contain "rc.1".
|
||||
//
|
||||
// This separation ensures forward compatibility: older readers can parse the
|
||||
// clean version field without errors, while newer readers can reconstruct the
|
||||
// full semantic version by combining version, prerelease, and build_metadata.
|
||||
//
|
||||
// If absent, the version field is used as-is.
|
||||
optional string prerelease = 3;
|
||||
// Optional semver build metadata.
|
||||
//
|
||||
// This field stores the build metadata portion of a semantic version separately
|
||||
// from the core version number. For example, if the full version is
|
||||
// "2.0.0-rc.1+build.123", the version field would contain "2.0.0", prerelease
|
||||
// would contain "rc.1", and build_metadata would contain "build.123".
|
||||
//
|
||||
// If absent, no build metadata is present.
|
||||
optional string build_metadata = 4;
|
||||
}
|
||||
|
||||
// The version of the writer that created this file.
|
||||
//
|
||||
// This information may be used to detect whether the file may have known bugs
|
||||
// associated with that writer.
|
||||
WriterVersion writer_version = 13;
|
||||
|
||||
// If present, the file position of the index metadata.
|
||||
optional uint64 index_section = 6;
|
||||
|
||||
// Version creation Timestamp, UTC timezone
|
||||
google.protobuf.Timestamp timestamp = 7;
|
||||
|
||||
// Optional version tag
|
||||
string tag = 8;
|
||||
|
||||
// Feature flags for readers.
|
||||
//
|
||||
// A bitmap of flags that indicate which features are required to be able to
|
||||
// read the table. If a reader does not recognize a flag that is set, it
|
||||
// should not attempt to read the dataset.
|
||||
//
|
||||
// Known flags:
|
||||
// * 1 << 0: deletion files are present
|
||||
// * 1 << 1: row ids are stable and stored as part of the fragment metadata.
|
||||
// * 1 << 2: use v2 format (deprecated)
|
||||
// * 1 << 3: table config is present
|
||||
// * 1 << 4: dataset uses multiple base paths
|
||||
// * 1 << 5: transaction file writes are disabled
|
||||
// * 1 << 6: data overlay files are present (see DataOverlayFile). Readers that do
|
||||
// not understand overlays must refuse the dataset, since ignoring an overlay
|
||||
// would silently return stale base values.
|
||||
uint64 reader_feature_flags = 9;
|
||||
|
||||
// Feature flags for writers.
|
||||
//
|
||||
// A bitmap of flags that indicate which features must be used when writing to the
|
||||
// dataset. If a writer does not recognize a flag that is set, it should not attempt to
|
||||
// write to the dataset.
|
||||
//
|
||||
// The flag identities are the same as for reader_feature_flags, but the values of
|
||||
// reader_feature_flags and writer_feature_flags are not required to be identical.
|
||||
uint64 writer_feature_flags = 10;
|
||||
|
||||
// The highest fragment ID that has been used so far.
|
||||
//
|
||||
// This ID is not guaranteed to be present in the current version, but it may
|
||||
// have been used in previous versions.
|
||||
//
|
||||
// For a single fragment, will be zero. For no fragments, will be absent.
|
||||
optional uint32 max_fragment_id = 11;
|
||||
|
||||
// Path to the transaction file, relative to `{root}/_transactions`. The file at that
|
||||
// location contains a wire-format serialized Transaction message representing the
|
||||
// transaction that created this version.
|
||||
//
|
||||
// This string field "transaction_file" may be empty if no transaction file was written.
|
||||
//
|
||||
// The path format is "{read_version}-{uuid}.txn" where {read_version} is the version of
|
||||
// the table the transaction read from (serialized to decimal with no padding digits),
|
||||
// and {uuid} is a hyphen-separated UUID.
|
||||
string transaction_file = 12;
|
||||
|
||||
// The file position of the transaction content. None if transaction is empty
|
||||
// This transaction content begins with the transaction content length as u32
|
||||
// If the transaction proto message has a length of `len`, the message ends at `len` + 4
|
||||
optional uint64 transaction_section = 21;
|
||||
|
||||
// The next unused row id. If zero, then the table does not have any rows.
|
||||
//
|
||||
// This is only used if the "stable_row_ids" feature flag is set.
|
||||
uint64 next_row_id = 14;
|
||||
|
||||
message DataStorageFormat {
|
||||
// The format of the data files (e.g. "lance")
|
||||
string file_format = 1;
|
||||
// The max format version of the data files. The format of the version can vary by
|
||||
// file_format and is not required to follow semver.
|
||||
//
|
||||
// Every file in this version of the dataset has the same file_format version.
|
||||
string version = 2;
|
||||
}
|
||||
|
||||
// The data storage format
|
||||
//
|
||||
// This specifies what format is used to store the data files.
|
||||
DataStorageFormat data_format = 15;
|
||||
|
||||
// Table config.
|
||||
//
|
||||
// Keys with the prefix "lance." are reserved for the Lance library. Other
|
||||
// libraries may wish to similarly prefix their configuration keys
|
||||
// appropriately.
|
||||
map<string, string> config = 16;
|
||||
|
||||
// Metadata associated with the table.
|
||||
//
|
||||
// This is a key-value map that can be used to store arbitrary metadata
|
||||
// associated with the table.
|
||||
//
|
||||
// This is different than configuration, which is used to tell libraries how
|
||||
// to read, write, or manage the table.
|
||||
//
|
||||
// This is different than schema metadata, which is used to describe the
|
||||
// data itself and is attached to the output schema of scans.
|
||||
map<string, string> table_metadata = 19;
|
||||
|
||||
// Field number 17 (`blob_dataset_version`) was used for a secondary blob dataset.
|
||||
reserved 17;
|
||||
reserved "blob_dataset_version";
|
||||
|
||||
// The base paths of data files.
|
||||
//
|
||||
// This is used to determine the base path of a data file. In common cases data file paths are under current dataset base path.
|
||||
// But for shallow cloning, importing file and other multi-tier storage cases, the actual data files could be outside of the current dataset.
|
||||
// This field is used with the `base_id` in `lance.file.File` and `lance.file.DeletionFile`.
|
||||
//
|
||||
// For example, if we have a dataset with base path `s3://bucket/dataset`, we have a DataFile with base_id 0, we get the actual data file path by:
|
||||
// base_paths[id = 0] + /data/ + file.path
|
||||
// the key(a.k.a index) starts from 0, increased by 1 for each new base path.
|
||||
repeated BasePath base_paths = 18;
|
||||
|
||||
// The branch of the dataset. None means main branch.
|
||||
optional string branch = 20;
|
||||
} // Manifest
|
||||
|
||||
// external dataset base path
|
||||
message BasePath {
|
||||
uint32 id = 1;
|
||||
// This is an alias name of the base path, it is optional.
|
||||
// When we use shallow clone and the target version is a tag, the tag name will be set here.
|
||||
optional string name = 2;
|
||||
// Flag indicating whether this path is a dataset root path or file directory:
|
||||
// - true: Path is a dataset root (actual files under subdirectories like `data`, '_deletions')
|
||||
// - false: Path is a direct file directory (scenario like importing files)
|
||||
bool is_dataset_root = 3;
|
||||
// Note: This absolute path will be directly used by Path:parse(),
|
||||
string path = 4;
|
||||
}
|
||||
|
||||
// Auxiliary Data attached to a version.
|
||||
// Only load on-demand.
|
||||
message VersionAuxData {
|
||||
// key-value metadata.
|
||||
map<string, bytes> metadata = 3;
|
||||
}
|
||||
|
||||
// Metadata describing an index.
|
||||
message IndexMetadata {
|
||||
// Unique ID of an index. It is unique across all the dataset versions.
|
||||
UUID uuid = 1;
|
||||
|
||||
// The columns to build the index. These refer to file.Field.id.
|
||||
repeated int32 fields = 2;
|
||||
|
||||
// Index name. Must be unique within one dataset version.
|
||||
string name = 3;
|
||||
|
||||
// The version of the dataset this index was built from.
|
||||
uint64 dataset_version = 4;
|
||||
|
||||
// A bitmap of the included fragment ids.
|
||||
//
|
||||
// This may by used to determine how much of the dataset is covered by the
|
||||
// index. This information can be retrieved from the dataset by looking at
|
||||
// the dataset at `dataset_version`. However, since the old version may be
|
||||
// deleted while the index is still in use, this information is also stored
|
||||
// in the index.
|
||||
//
|
||||
// The bitmap is stored as a 32-bit Roaring bitmap.
|
||||
bytes fragment_bitmap = 5;
|
||||
|
||||
// Details, specific to the index type, which are needed to load / interpret the index
|
||||
//
|
||||
// Indices should avoid putting large amounts of information in this field, as it will
|
||||
// bloat the manifest.
|
||||
//
|
||||
// Indexes are plugins, and so the format of the details message is flexible and not fully
|
||||
// defined by the table format. However, there are some conventions that should be followed:
|
||||
//
|
||||
// - When Lance APIs refer to indexes they will use the type URL of the index details as the
|
||||
// identifier for the index type. If a user provides a simple string identifier like
|
||||
// "btree" then it will be converted to "/lance.table.BTreeIndexDetails"
|
||||
// - Type URLs comparisons are case-insensitive. Thereform an index must have a unique type
|
||||
// URL ignoring case.
|
||||
google.protobuf.Any index_details = 6;
|
||||
|
||||
// The minimum lance version that this index is compatible with.
|
||||
optional int32 index_version = 7;
|
||||
|
||||
// Timestamp when the index was created (UTC timestamp in milliseconds since epoch)
|
||||
//
|
||||
// This field is optional for backward compatibility. For existing indices created before
|
||||
// this field was added, this will be None/null.
|
||||
optional uint64 created_at = 8;
|
||||
|
||||
// The base path index of the data file. Used when the file is imported or referred from another dataset.
|
||||
// Lance use it as key of the base_paths field in Manifest to determine the actual base path of the data file.
|
||||
optional uint32 base_id = 9;
|
||||
|
||||
// List of files and their sizes for this index segment.
|
||||
// This enables skipping HEAD calls when opening indices and allows reporting
|
||||
// of index sizes without extra IO.
|
||||
// If this is empty, the index files sizes are unknown.
|
||||
repeated IndexFile files = 10;
|
||||
}
|
||||
|
||||
// Metadata about a single file within an index segment.
|
||||
message IndexFile {
|
||||
// Path relative to the index directory (e.g., "index.idx", "auxiliary.idx")
|
||||
string path = 1;
|
||||
// Size of the file in bytes
|
||||
uint64 size_bytes = 2;
|
||||
}
|
||||
|
||||
// Index Section, containing a list of index metadata for one dataset version.
|
||||
message IndexSection {
|
||||
repeated IndexMetadata indices = 1;
|
||||
}
|
||||
|
||||
// A DataFragment is a set of files which represent the different columns of the same
|
||||
// rows. If column exists in the schema of a dataset, but the file for that column does
|
||||
// not exist within a DataFragment of that dataset, that column consists entirely of
|
||||
// nulls.
|
||||
message DataFragment {
|
||||
// The ID of a DataFragment is unique within a dataset.
|
||||
uint64 id = 1;
|
||||
|
||||
repeated DataFile files = 2;
|
||||
|
||||
// Optional overlay files for this fragment, which supply new values for a
|
||||
// subset of cells without rewriting the base data files. This MUST be empty
|
||||
// if the data overlay files feature flag (64) is not set in the manifest.
|
||||
//
|
||||
// Order is significant: a later entry is newer than an earlier one. When two
|
||||
// overlays cover the same (offset, field) and share a `committed_version`, the
|
||||
// later entry wins. See DataOverlayFile for the full resolution rules.
|
||||
repeated DataOverlayFile overlays = 11;
|
||||
|
||||
// File that indicates which rows, if any, should be considered deleted.
|
||||
DeletionFile deletion_file = 3;
|
||||
|
||||
// TODO: What's the simplest way we can allow an inline tombstone bitmap?
|
||||
|
||||
// A serialized RowIdSequence message (see rowids.proto).
|
||||
//
|
||||
// These are the row ids for the fragment, in order of the rows as they appear.
|
||||
// That is, if a fragment has 3 rows, and the row ids are [1, 42, 3], then the
|
||||
// first row is row 1, the second row is row 42, and the third row is row 3.
|
||||
oneof row_id_sequence {
|
||||
// If small (< 200KB), the row ids are stored inline.
|
||||
bytes inline_row_ids = 5;
|
||||
// Otherwise, stored as part of a file.
|
||||
ExternalFile external_row_ids = 6;
|
||||
} // row_id_sequence
|
||||
|
||||
oneof last_updated_at_version_sequence {
|
||||
// If small (< 200KB), the row latest updated versions are stored inline.
|
||||
bytes inline_last_updated_at_versions = 7;
|
||||
// Otherwise, stored as part of a file.
|
||||
ExternalFile external_last_updated_at_versions = 8;
|
||||
} // last_updated_at_version_sequence
|
||||
|
||||
oneof created_at_version_sequence {
|
||||
// If small (< 200KB), the row created at versions are stored inline.
|
||||
bytes inline_created_at_versions = 9;
|
||||
// Otherwise, stored as part of a file.
|
||||
ExternalFile external_created_at_versions = 10;
|
||||
} // created_at_version_sequence
|
||||
|
||||
// Number of original rows in the fragment, this includes rows that are now marked with
|
||||
// deletion tombstones. To compute the current number of rows, subtract
|
||||
// `deletion_file.num_deleted_rows` from this value.
|
||||
uint64 physical_rows = 4;
|
||||
}
|
||||
|
||||
message DataFile {
|
||||
// Path to the root relative to the dataset's URI.
|
||||
string path = 1;
|
||||
// The ids of the fields/columns in this file.
|
||||
//
|
||||
// When a DataFile object is created in memory, every value in fields is assigned -1 by
|
||||
// default. An object with a value in fields of -1 must not be stored to disk. -2 is
|
||||
// used for "tombstoned", meaning a field that is no longer in use. This is often
|
||||
// because the original field id was reassigned to a different data file.
|
||||
//
|
||||
// In Lance v1 IDs are assigned based on position in the file, offset by the max
|
||||
// existing field id in the table (if any already). So when a fragment is first created
|
||||
// with one file of N columns, the field ids will be 1, 2, ..., N. If a second fragment
|
||||
// is created with M columns, the field ids will be N+1, N+2, ..., N+M.
|
||||
//
|
||||
// In Lance v1 there is one field for each field in the input schema, this includes
|
||||
// nested fields (both struct and list). Fixed size list fields have only a single
|
||||
// field id (these are not considered nested fields in Lance v1).
|
||||
//
|
||||
// This allows column indices to be calculated from field IDs and the input schema.
|
||||
//
|
||||
// In Lance v2 the field IDs generally follow the same pattern but there is no
|
||||
// way to calculate the column index from the field ID. This is because a given
|
||||
// field could be encoded in many different ways, some of which occupy a different
|
||||
// number of columns. For example, a struct field could be encoded into N + 1 columns
|
||||
// or it could be encoded into a single packed column. To determine column indices
|
||||
// the column_indices property should be used instead.
|
||||
//
|
||||
// In Lance v1 these ids must be sorted but might not always be contiguous.
|
||||
repeated int32 fields = 2;
|
||||
// The top-level column indices for each field in the file.
|
||||
//
|
||||
// If the data file is version 1 then this property will be empty
|
||||
//
|
||||
// Otherwise there must be one entry for each field in `fields`.
|
||||
//
|
||||
// Some fields may not correspond to a top-level column in the file. In these cases
|
||||
// the index will -1.
|
||||
//
|
||||
// For example, consider the schema:
|
||||
//
|
||||
// - dimension: packed-struct (0):
|
||||
// - x: u32 (1)
|
||||
// - y: u32 (2)
|
||||
// - path: `list<u32>` (3)
|
||||
// - embedding: `fsl<768>` (4)
|
||||
// - fp64
|
||||
// - borders: `fsl<4>` (5)
|
||||
// - simple-struct (6)
|
||||
// - margin: fp64 (7)
|
||||
// - padding: fp64 (8)
|
||||
//
|
||||
// One possible column indices array could be:
|
||||
// [0, -1, -1, 1, 3, 4, 5, 6, 7]
|
||||
//
|
||||
// This reflects quite a few phenomenon:
|
||||
// - The packed struct is encoded into a single column and there is no top-level column
|
||||
// for the x or y fields
|
||||
// - The variable sized list is encoded into two columns
|
||||
// - The embedding is encoded into a single column (common for FSL of primitive) and there
|
||||
// is not "FSL column"
|
||||
// - The borders field actually does have an "FSL column"
|
||||
//
|
||||
// The column indices table may not have duplicates (other than -1)
|
||||
repeated int32 column_indices = 3;
|
||||
// The major file version used to create the file
|
||||
uint32 file_major_version = 4;
|
||||
// The minor file version used to create the file
|
||||
//
|
||||
// If both `file_major_version` and `file_minor_version` are set to 0,
|
||||
// then this is a version 0.1 or version 0.2 file.
|
||||
uint32 file_minor_version = 5;
|
||||
|
||||
// The known size of the file on disk in bytes.
|
||||
//
|
||||
// This is used to quickly find the footer of the file.
|
||||
//
|
||||
// When this is zero, it should be interpreted as "unknown".
|
||||
uint64 file_size_bytes = 6;
|
||||
|
||||
// The base path index of the data file. Used when the file is imported or referred from another dataset.
|
||||
// Lance use it as key of the base_paths field in Manifest to determine the actual base path of the data file.
|
||||
optional uint32 base_id = 7;
|
||||
} // DataFile
|
||||
|
||||
// An overlay file supplies new values for a subset of (row offset, field) cells
|
||||
// within a fragment, without rewriting the fragment's base data files. It is
|
||||
// used for efficient updates when only a small fraction of rows and/or columns
|
||||
// change.
|
||||
//
|
||||
// On read, a cell is resolved by consulting the fragment's overlays from newest
|
||||
// to oldest: the first overlay that covers that (offset, field) wins; if none
|
||||
// cover it, the value falls through to the base data file. Because deletions
|
||||
// take precedence over overlays, an overlay value for an offset that is also
|
||||
// marked deleted is dead and is ignored.
|
||||
//
|
||||
// The overlay's data file does NOT store a row-offset key column. Within a value
|
||||
// column, the position of a covered offset's value is the rank (0-based count of
|
||||
// set bits below it) of that offset within the field's coverage bitmap. Because
|
||||
// fields may cover different offset sets, the value columns of a single overlay
|
||||
// data file may have different lengths (which the Lance file format permits).
|
||||
message DataOverlayFile {
|
||||
// The data file storing the overlay's new cell values, one value column per
|
||||
// field in `data_file.fields`. No row-offset key column is stored.
|
||||
DataFile data_file = 1;
|
||||
|
||||
// Which (offset, field) cells this overlay provides values for.
|
||||
oneof coverage {
|
||||
// A single 32-bit Roaring bitmap of physical row offsets that applies to
|
||||
// every field in `data_file.fields` (a "dense" / rectangular overlay).
|
||||
// Every covered offset has a value for every field. This is the common case
|
||||
// for a plain UPDATE, where one SET list is applied to one set of rows.
|
||||
bytes shared_offset_bitmap = 2;
|
||||
// Per-field coverage for a "sparse" overlay, used when different fields cover
|
||||
// different offset sets (e.g. a MERGE with multiple WHEN MATCHED branches).
|
||||
FieldCoverage field_coverage = 4;
|
||||
}
|
||||
|
||||
// The dataset version at which this overlay became effective: the version of
|
||||
// the commit that introduced it, NOT the version it was read from. It is
|
||||
// stamped at commit time and re-stamped if the commit is retried, in the same
|
||||
// way as the created-at / last-updated-at version sequences.
|
||||
//
|
||||
// This drives two orderings:
|
||||
// * Versus index builds: an index whose `dataset_version` >= this value
|
||||
// already incorporates this overlay. Otherwise the overlay's covered cells
|
||||
// are excluded from index results for the affected fields and re-evaluated
|
||||
// against their current values (see the Data Overlay Files specification).
|
||||
// * Versus other overlays: when two overlays cover the same (offset, field),
|
||||
// the one with the higher `committed_version` wins. Overlays that share a
|
||||
// `committed_version` are ordered by their position in
|
||||
// `DataFragment.overlays`, where a later entry is newer and wins.
|
||||
uint64 committed_version = 3;
|
||||
}
|
||||
|
||||
// Per-field coverage for a sparse overlay.
|
||||
message FieldCoverage {
|
||||
// One entry per field in the overlay's `data_file.fields`, in the same order.
|
||||
// Each is a 32-bit Roaring bitmap of the physical row offsets covered for that
|
||||
// field. An offset present in a field's bitmap but mapped to a NULL value
|
||||
// means the cell is overridden to NULL (distinct from an offset that is absent,
|
||||
// which falls through to the base data file).
|
||||
repeated bytes offset_bitmaps = 1;
|
||||
}
|
||||
|
||||
// Deletion File
|
||||
//
|
||||
// The path of the deletion file is constructed as:
|
||||
// {root}/_deletions/{fragment_id}-{read_version}-{id}.{extension}
|
||||
// where {extension} depends on DeletionFileType.
|
||||
message DeletionFile {
|
||||
// Type of deletion file, intended as a way to increase efficiency of the storage of deleted row
|
||||
// offsets. If there are sparsely deleted rows, then ARROW_ARRAY is the most efficient. If there
|
||||
// are densely deleted rows, then BITMAP is the most efficient.
|
||||
enum DeletionFileType {
|
||||
// A single Int32Array of deleted row offsets, stored as an Arrow IPC file with one batch and
|
||||
// one column. Has a .arrow extension.
|
||||
ARROW_ARRAY = 0;
|
||||
// A Roaring Bitmap of deleted row offsets. Has a .bin extension.
|
||||
BITMAP = 1;
|
||||
}
|
||||
|
||||
// Type of deletion file.
|
||||
DeletionFileType file_type = 1;
|
||||
// The version of the dataset this deletion file was built from.
|
||||
uint64 read_version = 2;
|
||||
// An opaque id used to differentiate this file from others written by concurrent
|
||||
// writers.
|
||||
uint64 id = 3;
|
||||
// The number of rows that are marked as deleted.
|
||||
uint64 num_deleted_rows = 4;
|
||||
// The base path index of the deletion file. Used when the file is imported or referred from another
|
||||
// dataset. Lance uses it as key of the base_paths field in Manifest to determine the actual base
|
||||
// path of the deletion file.
|
||||
optional uint32 base_id = 7;
|
||||
} // DeletionFile
|
||||
|
||||
message ExternalFile {
|
||||
// Path to the file, relative to the root of the table.
|
||||
string path = 1;
|
||||
// The byte offset in the file where the data starts.
|
||||
uint64 offset = 2;
|
||||
// The size of the data in the file, in bytes.
|
||||
uint64 size = 3;
|
||||
}
|
||||
|
||||
// VectorIndexDetails and HnswParameters (formerly HnswIndexDetails) moved to index.proto
|
||||
|
||||
message FragmentReuseIndexDetails {
|
||||
|
||||
oneof content {
|
||||
// if < 200KB, store the content inline, otherwise store the InlineContent bytes in external file
|
||||
InlineContent inline = 1;
|
||||
ExternalFile external = 2;
|
||||
}
|
||||
|
||||
message InlineContent {
|
||||
repeated Version versions = 1;
|
||||
}
|
||||
|
||||
message FragmentDigest {
|
||||
uint64 id = 1;
|
||||
|
||||
uint64 physical_rows = 2;
|
||||
|
||||
uint64 num_deleted_rows = 3;
|
||||
}
|
||||
|
||||
// A summarized version of the RewriteGroup information in a Rewrite transaction
|
||||
message Group {
|
||||
// A roaring treemap of the changed row addresses.
|
||||
// When combined with the old fragment IDs and new fragment IDs,
|
||||
// it can recover the full mapping of old row addresses to either new row addresses or deleted.
|
||||
// this mapping can then be used to remap indexes or satisfy index queries for the new unindexed fragments.
|
||||
bytes changed_row_addrs = 1;
|
||||
|
||||
repeated FragmentDigest old_fragments = 2;
|
||||
|
||||
repeated FragmentDigest new_fragments = 3;
|
||||
}
|
||||
|
||||
message Version {
|
||||
// The dataset_version at the time the index adds this version entry
|
||||
uint64 dataset_version = 1;
|
||||
|
||||
repeated Group groups = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// MemWAL Index Types
|
||||
// ============================================================================
|
||||
|
||||
// Lifecycle status of a WAL shard. Drives drop-table two-phase commit:
|
||||
// a SEALED shard refuses new writer claims (reversible) until the drop
|
||||
// commits (the shard dir is deleted) or rolls back (status -> ACTIVE).
|
||||
enum ShardStatus {
|
||||
// Normal: the shard accepts writer claims.
|
||||
ACTIVE = 0;
|
||||
// A drop is in flight: claims are refused. Reversible to ACTIVE.
|
||||
SEALED = 1;
|
||||
}
|
||||
|
||||
// Shard manifest containing epoch-based fencing and WAL state.
|
||||
// Each shard has exactly one active writer at any time.
|
||||
message ShardManifest {
|
||||
// Shard identifier (UUID v4).
|
||||
UUID shard_id = 11;
|
||||
|
||||
// Manifest version number.
|
||||
// Matches the version encoded in the filename.
|
||||
uint64 version = 1;
|
||||
|
||||
// Shard spec ID this shard was created with.
|
||||
// Set at shard creation and immutable thereafter.
|
||||
// A value of 0 indicates a manually-created shard not governed by any spec.
|
||||
uint32 shard_spec_id = 10;
|
||||
|
||||
// Computed shard field values as raw Arrow scalar bytes, keyed by shard
|
||||
// field id. The byte encoding follows Arrow's little-endian convention:
|
||||
// int32 is 4 LE bytes, utf8 is raw UTF-8 bytes, etc. The receiver looks
|
||||
// up the result_type from the ShardingSpec to interpret each value.
|
||||
repeated ShardFieldEntry shard_field_entries = 14;
|
||||
|
||||
// Writer fencing token - monotonically increasing.
|
||||
// A writer must increment this when claiming the shard.
|
||||
uint64 writer_epoch = 2;
|
||||
|
||||
// The most recent WAL entry position that has been flushed to a MemTable.
|
||||
// During recovery, replay starts from replay_after_wal_entry_position + 1.
|
||||
// WAL positions are 1-based, so the default value 0 unambiguously means
|
||||
// "no flush has ever stamped this shard" and recovery replays from 1.
|
||||
uint64 replay_after_wal_entry_position = 3;
|
||||
|
||||
// The most recent WAL entry position observed at the time the manifest was
|
||||
// updated. WAL positions are 1-based; default 0 means no entry has been
|
||||
// written yet. This is a hint, not authoritative - recovery must list
|
||||
// files to find actual state.
|
||||
uint64 wal_entry_position_last_seen = 4;
|
||||
|
||||
// Generation to assign to the next SSTable (incremented after each MemTable flush).
|
||||
uint64 current_generation = 6;
|
||||
|
||||
// Field 7 removed: compaction progress lives in
|
||||
// MemWalIndexDetails.compacted_sstables.
|
||||
|
||||
// List of SSTables created by flushing MemTables and their directory paths.
|
||||
repeated SsTable sstables = 8;
|
||||
|
||||
// Lifecycle status. Default ACTIVE; SEALED marks an in-flight drop
|
||||
// (drop-table 2PC). A SEALED manifest refuses claims at claim_epoch.
|
||||
ShardStatus status = 15;
|
||||
}
|
||||
|
||||
// A shard field value stored as raw Arrow scalar bytes.
|
||||
message ShardFieldEntry {
|
||||
// Shard field id (matches ShardingField.field_id in the ShardingSpec).
|
||||
string field_id = 1;
|
||||
|
||||
// Raw Arrow scalar value bytes in little-endian encoding.
|
||||
// The data type is determined by the result_type of the matching ShardingField.
|
||||
bytes value = 2;
|
||||
}
|
||||
|
||||
// An SSTable: the immutable result of flushing a MemTable, stored as a Lance dataset.
|
||||
message SsTable {
|
||||
// Generation number identifying this SSTable.
|
||||
uint64 generation = 1;
|
||||
|
||||
// Directory name relative to the shard directory.
|
||||
string path = 2;
|
||||
}
|
||||
|
||||
// A pointer to the latest SSTable compacted for a shard.
|
||||
message CompactedSsTable {
|
||||
// Shard identifier (UUID v4).
|
||||
UUID shard_id = 1;
|
||||
|
||||
// Generation of the latest SSTable compacted into the base table for this shard.
|
||||
uint64 generation = 2;
|
||||
}
|
||||
|
||||
// Tracks which compacted SSTable generation a base table index has been rebuilt to cover.
|
||||
// Used to determine whether to read from SSTable indexes or base table.
|
||||
message IndexCatchupProgress {
|
||||
// Name of the base table index (must match an entry in maintained_indexes).
|
||||
string index_name = 1;
|
||||
|
||||
// Per-shard progress: the generation up to which this index covers.
|
||||
// If a shard is not present, the index is assumed to be fully caught up
|
||||
// (i.e., caught_up_generation >= compacted_generation for that shard).
|
||||
repeated CompactedSsTable caught_up_generations = 2;
|
||||
}
|
||||
|
||||
// Index details for MemWAL Index, stored in IndexMetadata.index_details.
|
||||
// This is the centralized structure for all MemWAL metadata:
|
||||
// - Configuration (sharding specs, indexes to maintain)
|
||||
// - SSTable compaction progress
|
||||
// - Shard state snapshots
|
||||
//
|
||||
// Writers read this index to get configuration before writing.
|
||||
// Readers may use shard snapshots in this index as a point-in-time
|
||||
// optimization. Readers that need the latest shard set should list shard
|
||||
// directories in storage and read each shard's latest manifest.
|
||||
// A background process updates the index periodically to keep shard snapshots current.
|
||||
//
|
||||
// Shard snapshots are stored as a Lance file with one row per shard.
|
||||
// The schema records shard discovery fields. Full mutable shard state remains
|
||||
// authoritative in the shard manifest files.
|
||||
// shard_id: utf8
|
||||
// shard_spec_id: uint32
|
||||
// shard_field_{field_id}: typed per the matching ShardingField.result_type
|
||||
message MemWalIndexDetails {
|
||||
// Snapshot timestamp (Unix timestamp in milliseconds).
|
||||
int64 snapshot_ts_millis = 1;
|
||||
|
||||
// Number of shards in the snapshot.
|
||||
// Used to determine storage format without reading the snapshot data.
|
||||
uint32 num_shards = 2;
|
||||
|
||||
// Inline shard snapshots for small shard counts.
|
||||
// When num_shards <= threshold (implementation-defined, e.g., 100),
|
||||
// snapshots are stored inline as serialized bytes.
|
||||
// Format: Lance file bytes with the shard snapshot schema.
|
||||
optional bytes inline_snapshots = 3;
|
||||
|
||||
// Sharding specs defining how to derive shard identifiers.
|
||||
// This configuration determines how rows are partitioned into shards.
|
||||
repeated ShardingSpec sharding_specs = 7;
|
||||
|
||||
// Indexes from the base table to maintain in MemTables.
|
||||
// These are index names referencing indexes defined on the base table.
|
||||
// The primary key btree index is always maintained implicitly and
|
||||
// should not be listed here.
|
||||
//
|
||||
// For vector indexes, MemTables inherit quantization parameters (PQ codebook,
|
||||
// SQ params) from the base table index to ensure distance comparability.
|
||||
repeated string maintained_indexes = 8;
|
||||
|
||||
// Latest SSTable compacted into the base table for each shard.
|
||||
// This is updated atomically with merge-insert data commits, enabling
|
||||
// conflict resolution when multiple compactors operate concurrently.
|
||||
//
|
||||
// Note: This is separate from shard snapshots because:
|
||||
// 1. compacted_sstables is updated by compactors (atomic with data commit)
|
||||
// 2. shard snapshots are updated by background index builder
|
||||
repeated CompactedSsTable compacted_sstables = 9;
|
||||
|
||||
// Per-index catchup progress tracking.
|
||||
// When data is compacted into the base table, base table indexes are rebuilt
|
||||
// asynchronously. This field tracks which generation each index covers.
|
||||
//
|
||||
// For indexed queries, if an index's caught_up_generation < compacted_generation,
|
||||
// readers should use SSTable indexes for the gap instead of
|
||||
// scanning unindexed data in the base table.
|
||||
//
|
||||
// If an index is not present in this list, it is assumed to be fully caught up.
|
||||
repeated IndexCatchupProgress index_catchup = 10;
|
||||
|
||||
// Default ShardWriter configuration values for this MemWAL index.
|
||||
//
|
||||
// A free-form string map persisted so that every writer — across
|
||||
// processes and restarts — starts from the same default writer
|
||||
// configuration. These are defaults only: an individual writer may
|
||||
// still override any value at runtime in its own ShardWriterConfig
|
||||
// (which is not persisted).
|
||||
map<string, string> writer_config_defaults = 11;
|
||||
}
|
||||
|
||||
// Sharding spec definition.
|
||||
message ShardingSpec {
|
||||
// Unique identifier for this spec within the index.
|
||||
// IDs are never reused.
|
||||
uint32 spec_id = 1;
|
||||
|
||||
// Sharding field definitions that determine how to compute shard identifiers.
|
||||
repeated ShardingField fields = 2;
|
||||
}
|
||||
|
||||
// Sharding field definition.
|
||||
message ShardingField {
|
||||
// Unique string identifier for this shard field.
|
||||
string field_id = 1;
|
||||
|
||||
// Field IDs referencing source columns in the schema.
|
||||
repeated int32 source_ids = 2;
|
||||
|
||||
// Well-known shard transform name (e.g., "identity", "year", "bucket").
|
||||
// Mutually exclusive with expression.
|
||||
optional string transform = 3;
|
||||
|
||||
// DataFusion SQL expression for custom logic.
|
||||
// Mutually exclusive with transform.
|
||||
optional string expression = 4;
|
||||
|
||||
// Output type of the shard value (Arrow type name).
|
||||
string result_type = 5;
|
||||
|
||||
// Transform parameters (e.g., num_buckets for bucket transform).
|
||||
map<string, string> parameters = 6;
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package lance.datafusion;
|
||||
|
||||
// Identifies a Lance dataset for remote reconstruction.
|
||||
//
|
||||
// Two modes:
|
||||
// 1. uri + serialized_manifest (fast): remote executor skips manifest read.
|
||||
// 2. uri + version + etag (lightweight): remote executor loads manifest from storage.
|
||||
message TableIdentifier {
|
||||
string uri = 1;
|
||||
uint64 version = 2;
|
||||
optional string manifest_etag = 3;
|
||||
optional bytes serialized_manifest = 4;
|
||||
map<string, string> storage_options = 5;
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
import "file.proto";
|
||||
import "table.proto";
|
||||
import "google/protobuf/any.proto";
|
||||
|
||||
package lance.table;
|
||||
|
||||
// A transaction represents the changes to a dataset.
|
||||
//
|
||||
// This has two purposes:
|
||||
// 1. When retrying a commit, the transaction can be used to re-build an updated
|
||||
// manifest.
|
||||
// 2. When there's a conflict, this can be used to determine whether the other
|
||||
// transaction is compatible with this one.
|
||||
message Transaction {
|
||||
// The version of the dataset this transaction was built from.
|
||||
//
|
||||
// For example, for a delete transaction this means the version of the dataset
|
||||
// that was read from while evaluating the deletion predicate.
|
||||
uint64 read_version = 1;
|
||||
|
||||
// The UUID that unique identifies a transaction.
|
||||
string uuid = 2;
|
||||
|
||||
// Optional version tag.
|
||||
string tag = 3;
|
||||
|
||||
// Optional properties for the transaction
|
||||
// __lance_commit_message is a reserved key
|
||||
map<string, string> transaction_properties = 4;
|
||||
|
||||
// Add new rows to the dataset.
|
||||
message Append {
|
||||
// The new fragments to append.
|
||||
//
|
||||
// Fragment IDs are not yet assigned.
|
||||
repeated DataFragment fragments = 1;
|
||||
}
|
||||
|
||||
// Mark rows as deleted.
|
||||
message Delete {
|
||||
// The fragments to update
|
||||
//
|
||||
// The fragment IDs will match existing fragments in the dataset.
|
||||
repeated DataFragment updated_fragments = 1;
|
||||
// The fragments to delete entirely.
|
||||
repeated uint64 deleted_fragment_ids = 2;
|
||||
// The predicate that was evaluated
|
||||
//
|
||||
// This may be used to determine whether the delete would have affected
|
||||
// files written by a concurrent transaction.
|
||||
string predicate = 3;
|
||||
}
|
||||
|
||||
// Create or overwrite the entire dataset.
|
||||
message Overwrite {
|
||||
// The new fragments
|
||||
//
|
||||
// Fragment IDs are not yet assigned.
|
||||
repeated DataFragment fragments = 1;
|
||||
// The new schema
|
||||
repeated lance.file.Field schema = 2;
|
||||
// Schema metadata.
|
||||
map<string, bytes> schema_metadata = 3;
|
||||
// Key-value pairs to merge with existing config.
|
||||
map<string, string> config_upsert_values = 4;
|
||||
// The base paths to be added for the initial dataset creation
|
||||
repeated BasePath initial_bases = 5;
|
||||
}
|
||||
|
||||
// Add or replace a new secondary index.
|
||||
//
|
||||
// This is also used to remove an index (we are replacing it with nothing)
|
||||
//
|
||||
// - new_indices: the modified indices, empty if dropping indices only
|
||||
// - removed_indices: the indices that are being replaced
|
||||
message CreateIndex {
|
||||
repeated IndexMetadata new_indices = 1;
|
||||
repeated IndexMetadata removed_indices = 2;
|
||||
}
|
||||
|
||||
// An operation that rewrites but does not change the data in the table. These
|
||||
// kinds of operations just rearrange data.
|
||||
message Rewrite {
|
||||
// The old fragments that are being replaced
|
||||
//
|
||||
// DEPRECATED: use groups instead.
|
||||
//
|
||||
// These should all have existing fragment IDs.
|
||||
repeated DataFragment old_fragments = 1;
|
||||
// The new fragments
|
||||
//
|
||||
// DEPRECATED: use groups instead.
|
||||
//
|
||||
// These fragments IDs are not yet assigned.
|
||||
repeated DataFragment new_fragments = 2;
|
||||
|
||||
// During a rewrite an index may be rewritten. We only serialize the UUID
|
||||
// since a rewrite should not change the other index parameters.
|
||||
message RewrittenIndex {
|
||||
// The id of the index that will be replaced
|
||||
UUID old_id = 1;
|
||||
// the id of the new index
|
||||
UUID new_id = 2;
|
||||
// the new index details
|
||||
google.protobuf.Any new_index_details = 3;
|
||||
// the version of the new index
|
||||
uint32 new_index_version = 4;
|
||||
// Files in the new index with their sizes.
|
||||
// Empty if file sizes are not available (e.g. older writers).
|
||||
repeated IndexFile new_index_files = 5;
|
||||
}
|
||||
|
||||
// A group of rewrite files that are all part of the same rewrite.
|
||||
message RewriteGroup {
|
||||
// The old fragment that is being replaced
|
||||
//
|
||||
// This should have an existing fragment ID.
|
||||
repeated DataFragment old_fragments = 1;
|
||||
// The new fragment
|
||||
//
|
||||
// The ID should have been reserved by an earlier
|
||||
// reserve operation
|
||||
repeated DataFragment new_fragments = 2;
|
||||
}
|
||||
|
||||
// Groups of files that have been rewritten
|
||||
repeated RewriteGroup groups = 3;
|
||||
// Indices that have been rewritten
|
||||
repeated RewrittenIndex rewritten_indices = 4;
|
||||
}
|
||||
|
||||
// An operation that merges in a new column, altering the schema.
|
||||
message Merge {
|
||||
// The updated fragments
|
||||
//
|
||||
// These should all have existing fragment IDs.
|
||||
repeated DataFragment fragments = 1;
|
||||
// The new schema
|
||||
repeated lance.file.Field schema = 2;
|
||||
// Schema metadata.
|
||||
map<string, bytes> schema_metadata = 3;
|
||||
}
|
||||
|
||||
// An operation that projects a subset of columns, altering the schema.
|
||||
message Project {
|
||||
// The new schema
|
||||
repeated lance.file.Field schema = 1;
|
||||
}
|
||||
|
||||
// An operation that restores a dataset to a previous version.
|
||||
message Restore {
|
||||
// The version to restore to
|
||||
uint64 version = 1;
|
||||
}
|
||||
|
||||
// An operation that reserves fragment ids for future use in
|
||||
// a rewrite operation.
|
||||
message ReserveFragments {
|
||||
uint32 num_fragments = 1;
|
||||
}
|
||||
|
||||
// An operation that clones a dataset.
|
||||
message Clone {
|
||||
// - true: Performs a metadata-only clone (copies manifest without data files).
|
||||
// The cloned dataset references original data through `base_paths`,
|
||||
// suitable for experimental scenarios or rapid metadata migration.
|
||||
// - false: Performs a full deep clone using the underlying object storage's native
|
||||
// copy API (e.g., S3 CopyObject, GCS rewrite). This leverages server-side
|
||||
// bulk copy operations to bypass download/upload bottlenecks, achieving
|
||||
// near-linear speedup for large datasets (typically 3-10x faster than
|
||||
// manual file transfers). The operation maintains atomicity and data
|
||||
// integrity guarantees provided by the storage backend.
|
||||
bool is_shallow = 1;
|
||||
// the reference name in the source dataset
|
||||
// in most cases it should be the branch or tag name in the source dataset
|
||||
optional string ref_name = 2;
|
||||
// the version of the source dataset for cloning
|
||||
uint64 ref_version = 3;
|
||||
// the absolute base path of the source dataset for cloning
|
||||
string ref_path = 4;
|
||||
// if the target dataset is a branch, this is the branch name of the target dataset
|
||||
optional string branch_name = 5;
|
||||
}
|
||||
|
||||
// Exact set of key hashes for conflict detection.
|
||||
// Used when the number of inserted rows is small.
|
||||
message ExactKeySetFilter {
|
||||
// 64-bit hashes of the inserted row keys.
|
||||
repeated uint64 key_hashes = 1;
|
||||
}
|
||||
|
||||
// Bloom filter for key existence tests.
|
||||
// Used when the number of rows is large.
|
||||
message BloomFilter {
|
||||
// Bitset backing the bloom filter (SBBF format).
|
||||
bytes bitmap = 1;
|
||||
// Number of bits in the bitmap.
|
||||
uint32 num_bits = 2;
|
||||
// Number of items the filter was sized for.
|
||||
// Used for intersection validation (filters with different sizes cannot be compared).
|
||||
// Default: 8192
|
||||
uint64 number_of_items = 3;
|
||||
// False positive probability the filter was sized for.
|
||||
// Used for intersection validation (filters with different parameters cannot be compared).
|
||||
// Default: 0.00057
|
||||
double probability = 4;
|
||||
}
|
||||
|
||||
// A filter for checking key existence in set of rows inserted by a merge insert operation.
|
||||
// Only created when the merge insert's ON columns match the schema's unenforced primary key.
|
||||
// The presence of this filter indicates strict primary key conflict detection should be used.
|
||||
// Can use either an exact set (for small row counts) or a Bloom filter (for large row counts).
|
||||
message KeyExistenceFilter {
|
||||
// Field IDs of columns participating in the key (must match unenforced primary key).
|
||||
repeated int32 field_ids = 1;
|
||||
// The underlying data structure storing the key hashes.
|
||||
oneof data {
|
||||
// Exact set of key hashes (used for small number of rows).
|
||||
ExactKeySetFilter exact = 2;
|
||||
// Bloom filter (used for large number of rows).
|
||||
BloomFilter bloom = 3;
|
||||
}
|
||||
}
|
||||
|
||||
// Serialized as sorted distinct local physical row offsets within the fragment (0-based).
|
||||
message UInt32List {
|
||||
repeated uint32 values = 1;
|
||||
}
|
||||
|
||||
// An operation that updates rows but does not add or remove rows.
|
||||
message Update {
|
||||
// The fragments that have been removed. These are fragments where all rows
|
||||
// have been updated and moved to a new fragment.
|
||||
repeated uint64 removed_fragment_ids = 1;
|
||||
// The fragments that have been updated.
|
||||
repeated DataFragment updated_fragments = 2;
|
||||
// The new fragments where updated rows have been moved to.
|
||||
repeated DataFragment new_fragments = 3;
|
||||
// The ids of the fields that have been modified.
|
||||
repeated uint32 fields_modified = 4;
|
||||
/// SSTables to mark as compacted after this transaction.
|
||||
repeated CompactedSsTable compacted_sstables = 5;
|
||||
/// The fields that used to judge whether to preserve the new frag's id into
|
||||
/// the frag bitmap of the specified indices.
|
||||
repeated uint32 fields_for_preserving_frag_bitmap = 6;
|
||||
// The mode of update
|
||||
UpdateMode update_mode = 7;
|
||||
// Filter for checking existence of keys in newly inserted rows, used for conflict detection.
|
||||
// Only tracks keys from INSERT operations during merge insert, not updates.
|
||||
optional KeyExistenceFilter inserted_rows = 8;
|
||||
// Per-fragment physical row offsets that matched an update_columns hash join (RewriteColumns).
|
||||
map<uint64, UInt32List> updated_fragment_offsets = 9;
|
||||
}
|
||||
|
||||
// The mode of update operation
|
||||
enum UpdateMode {
|
||||
|
||||
/// rows are deleted in current fragments and rewritten in new fragments.
|
||||
/// This is most optimal when the majority of columns are being rewritten
|
||||
/// or only a few rows are being updated.
|
||||
REWRITE_ROWS = 0;
|
||||
|
||||
/// within each fragment, columns are fully rewritten and inserted as new data files.
|
||||
/// Old versions of columns are tombstoned. This is most optimal when most rows are affected
|
||||
/// but a small subset of columns are affected.
|
||||
REWRITE_COLUMNS = 1;
|
||||
}
|
||||
|
||||
// An entry for a map update. If value is not set, the key will be removed from the map.
|
||||
message UpdateMapEntry {
|
||||
// The key of the map entry to update.
|
||||
string key = 1;
|
||||
// The value to set for the key.
|
||||
optional string value = 2;
|
||||
}
|
||||
|
||||
message UpdateMap {
|
||||
repeated UpdateMapEntry update_entries = 1;
|
||||
// If true, the map will be replaced entirely with the new entries.
|
||||
// If false, the new entries will be merged with the existing map.
|
||||
bool replace = 2;
|
||||
}
|
||||
|
||||
// An operation that updates the table config, table metadata, schema metadata,
|
||||
// or field metadata.
|
||||
message UpdateConfig {
|
||||
UpdateMap config_updates = 6;
|
||||
UpdateMap table_metadata_updates = 7;
|
||||
UpdateMap schema_metadata_updates = 8;
|
||||
map<int32, UpdateMap> field_metadata_updates = 9;
|
||||
|
||||
// Deprecated -------------------------------
|
||||
map<string, string> upsert_values = 1;
|
||||
repeated string delete_keys = 2;
|
||||
map<string, string> schema_metadata = 3;
|
||||
map<uint32, FieldMetadataUpdate> field_metadata = 4;
|
||||
|
||||
message FieldMetadataUpdate {
|
||||
map<string, string> metadata = 5;
|
||||
}
|
||||
}
|
||||
|
||||
message DataReplacementGroup {
|
||||
uint64 fragment_id = 1;
|
||||
DataFile new_file = 2;
|
||||
}
|
||||
|
||||
// An operation that replaces the data in a region of the table with new data.
|
||||
message DataReplacement {
|
||||
repeated DataReplacementGroup replacements = 1;
|
||||
}
|
||||
|
||||
// Overlay files to append to a single fragment, in order (the last entry is
|
||||
// newest). The overlays are appended to the fragment's existing `overlays`
|
||||
// list; they do not replace it, so overlays written by concurrent commits are
|
||||
// preserved.
|
||||
message DataOverlayGroup {
|
||||
uint64 fragment_id = 1;
|
||||
// Each DataOverlayFile.committed_version is left 0 by the writer and stamped
|
||||
// to the new dataset version at commit time (re-stamped on retry), in the
|
||||
// same way as the created-at / last-updated-at version sequences. The fields
|
||||
// touched are read from each overlay's `data_file.fields`.
|
||||
repeated DataOverlayFile overlays = 2;
|
||||
}
|
||||
|
||||
// Attach overlay files to fragments, supplying new values for a subset of
|
||||
// (row offset, field) cells without rewriting the fragments' base data files.
|
||||
// See the DataOverlayFile message in table.proto for resolution, coverage, and
|
||||
// versioning rules, and the Data Overlay Files and Transactions specifications
|
||||
// for the (intentionally permissive) conflict semantics.
|
||||
message DataOverlay {
|
||||
repeated DataOverlayGroup groups = 1;
|
||||
}
|
||||
|
||||
// Update SSTable compaction progress in the MemWAL index.
|
||||
// This operation is used during merge-insert to atomically record which
|
||||
// SSTables have been compacted into the base table.
|
||||
message UpdateMemWalState {
|
||||
// SSTables being marked as compacted.
|
||||
repeated CompactedSsTable compacted_sstables = 1;
|
||||
}
|
||||
|
||||
// An operation that updates base paths in the dataset.
|
||||
message UpdateBases {
|
||||
// The new base paths to add to the manifest.
|
||||
repeated BasePath new_bases = 1;
|
||||
}
|
||||
|
||||
// The operation of this transaction.
|
||||
oneof operation {
|
||||
Append append = 100;
|
||||
Delete delete = 101;
|
||||
Overwrite overwrite = 102;
|
||||
CreateIndex create_index = 103;
|
||||
Rewrite rewrite = 104;
|
||||
Merge merge = 105;
|
||||
Restore restore = 106;
|
||||
ReserveFragments reserve_fragments = 107;
|
||||
Update update = 108;
|
||||
Project project = 109;
|
||||
UpdateConfig update_config = 110;
|
||||
DataReplacement data_replacement = 111;
|
||||
UpdateMemWalState update_mem_wal_state = 112;
|
||||
Clone clone = 113;
|
||||
UpdateBases update_bases = 114;
|
||||
DataOverlay data_overlay = 115;
|
||||
}
|
||||
|
||||
// Fields 200/202 (`blob_append` / `blob_overwrite`) previously represented blob dataset ops.
|
||||
reserved 200, 202;
|
||||
reserved "blob_append", "blob_overwrite";
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
.env
|
||||
@@ -0,0 +1,90 @@
|
||||
# Rust Guidelines
|
||||
|
||||
Also see [root AGENTS.md](../AGENTS.md) for cross-language standards.
|
||||
|
||||
## Code Style
|
||||
|
||||
- Use `Vec::with_capacity()` when size is known or estimable — prefer over-estimating capacity to multiple reallocations.
|
||||
- Wrap large or expensive-to-clone struct fields (maps, protobuf metadata, schemas) in `Arc<T>` to avoid deep copies.
|
||||
- Use `Box::pin(...)` or `.boxed()` but never both — `.boxed()` already returns `Pin<Box<...>>`.
|
||||
- Remove dead code instead of adding `#[allow(dead_code)]`. Delete unused constants instead of reducing visibility.
|
||||
- Use `column_by_name()` for `RecordBatch` column access in production code; use `batch["column_name"]` in tests.
|
||||
- Use `PrimitiveArray::<T>::from(vec)` (zero-copy) instead of `from_iter_values(vec)` for Vec-to-PrimitiveArray conversion.
|
||||
- Implement `Default` trait on config/options structs instead of standalone `default_*()` helpers.
|
||||
- Place `#[cfg(test)] mod tests` as a single block at the bottom of each file — no production code after it.
|
||||
- Place `use` imports at the top of the file, not inline within function bodies.
|
||||
- Extract substantial new logic (bin packing, scheduling) into dedicated submodules instead of inlining into large files.
|
||||
- Delete obsolete internal (`pub(crate)` / private) methods in the same PR that introduces their replacements. For public API methods, follow the deprecation path in root AGENTS.md instead.
|
||||
- Choose log levels by audience: `debug!` for routine/high-frequency ops, `info!` for infrequent operator-visible state changes, `warn!` for unexpected conditions.
|
||||
|
||||
## Concurrency
|
||||
|
||||
- The closure passed to `spawn_cpu()` must only consume CPU and return — it must **never** wait on anything: **no channels** (blocking send/recv), **no I/O**, **no locks**, and no `block_on`/`.blocking_*`. The CPU pool can collapse to a single worker in resource-constrained environments (`<= 3` CPUs), so a parked closure can deadlock the whole pool with a silent 0% hang. Keep the waiting in surrounding async code and hand only the pure-CPU work to `spawn_cpu()`. Only dispatch substantial work (rule of thumb: ~100µs+ of CPU); below that the pool overhead outweighs the benefit and the work is better left inline. See the doc comment on `spawn_cpu` for the rationale.
|
||||
|
||||
## API Design
|
||||
|
||||
- Use `with_`-prefixed builder methods for optional config (e.g., `MyStruct::new(required).with_option(v)`) — don't create separate constructor variants.
|
||||
- For public APIs, prefer `Into<T>` or `AsRef<T>` trait bounds for flexible inputs.
|
||||
- Prefer `pub(crate)` over `pub` for crate-internal items. Use `pub use` re-exports for the actual public API surface.
|
||||
- Use enums instead of magic numbers for format versions, variant types, and discriminators — leverage exhaustive `match`.
|
||||
- Use strongly-typed structs instead of `HashMap<String, String>` in APIs — convert to strings only at serialization boundaries.
|
||||
- Keep `RowAddr` (physical fragment+offset) and `RowId` (stable logical identifier) as distinct types — never raw `u64` for both.
|
||||
- Use `RowAddress` from `lance-core/src/utils/address.rs` instead of raw bitwise operations on row addresses.
|
||||
- Use `RowAddrTreeMap`/`RoaringBitmap` instead of `Vec<Range<u64>>` for physical row selections.
|
||||
- Use logical row counts (`num_rows()`) instead of `physical_rows` for user-facing metrics — subtract deletions.
|
||||
- Keep traits minimal — only core abstraction methods. Move helpers to standalone functions and config to struct fields.
|
||||
- Get column/field types from schema metadata — never materialize data rows just to inspect types.
|
||||
- Use stable, versioned serialization formats for persistent storage (e.g., index files) — avoid unstable cross-version formats.
|
||||
- Use Arrow's type-safe access (`ArrayAccessor` trait bounds, `as_*_array` helpers) instead of `arrow::compute::cast` + `downcast_ref`. Prefer `_opt` variants (e.g., `as_string_opt`) unless the data type has already been verified.
|
||||
- In `lance-io/`, use single-syscall writes for local filesystem I/O — don't reuse cloud multipart upload machinery.
|
||||
|
||||
## Error Handling
|
||||
|
||||
- Never use `.unwrap()`, `.expect()`, `panic!()`, or `assert!()` in library code for fallible operations — use `?` with `Result` and proper error types. Reserve `.unwrap()` for tests only.
|
||||
- Avoid bare `.unwrap()`; use `if let`, `match`, `let ... else`, `?`, or combinators. Never `.is_none()` followed by `.unwrap()`. If unavoidable, use `.expect("reason")`.
|
||||
- Return `LanceError::NotSupported` instead of `todo!()` or `unimplemented!()` for unsupported code paths. Test with `Result::Err` assertions, not `#[should_panic]`.
|
||||
- Match `Error` variant to root cause: `Error::invalid_input` for caller data issues, `Error::corrupt_file` for format/integrity issues, `Error::not_found` for missing resources, `Error::io` for I/O failures.
|
||||
- Include full context in error messages — variable names, values, sizes, types, indices. Not generic messages like `"Invalid chunk size"`.
|
||||
- Use `checked_add`/`checked_mul` instead of `wrapping_add`/`wrapping_mul` for counters and IDs — return an error on overflow.
|
||||
- Prefer `debug_assert!` over `assert!` for non-safety invariants; reserve `assert!` for conditions preventing data corruption. Always include descriptive messages.
|
||||
- Don't silently guard against impossible conditions — use `debug_assert!`, return an explicit error, or remove the check.
|
||||
- Log warnings on best-effort/cleanup failures instead of silently swallowing or propagating errors.
|
||||
- Log warnings for silent no-ops (skipped operations); omit warnings before errors since the error message is sufficient.
|
||||
- Avoid `unwrap_or(default)` on map lookups for required config params — use `.ok_or_else(|| Error::...)` and verify key names match between serialization and deserialization.
|
||||
- Advance all parallel iterators before any `continue` branches — early exits that skip `.next()` calls cause misalignment.
|
||||
- Bind `iter.next()` with `let Some(x) = iter.next() else { ... }` — never call `.next()` twice to check-then-use.
|
||||
|
||||
## Naming
|
||||
|
||||
- Reserve `_`-prefixed names for truly unused bindings — if a variable is read, drop the underscore.
|
||||
- Prefix boolean variables with `is_` or `has_` instead of ambiguous `with_` or bare adjectives.
|
||||
- Name booleans so `false` (zero/`Default::default()`) is the desired default — use `disable_*` instead of `enable_*` when the feature should be on by default.
|
||||
- Name functions to match their actual scope — e.g., `handle_partition_system_columns` not `handle_system_columns` if only a subset is handled.
|
||||
|
||||
## Testing
|
||||
|
||||
- Use `record_batch!()` from `arrow_array` to construct `RecordBatch` in tests instead of manual Schema/Arc/try_new boilerplate.
|
||||
- Use `gen_batch()` builder API (`.col()`, `.into_reader_rows()`) for test data setup instead of manual Arrow construction.
|
||||
- Use `.try_into_batch()` instead of `.try_into_stream().try_collect()` for scanner results in tests.
|
||||
- Use plain `"memory://"` URIs in tests — no atomic counters or unique suffixes needed.
|
||||
- Assert on both error variant (`assert!(matches!(error, ErrorType::Variant { .. }))`) and message content — don't just check `is_err()`.
|
||||
|
||||
## Documentation
|
||||
|
||||
- Add doc comments to public API elements that convey semantic meaning, valid values, and effects — don't restate type signatures.
|
||||
- Document enum variant doc comments with behavioral semantics, not just labels. For numeric parameters, state whether it's an id, count, index, etc.
|
||||
- Add doc comments to magic constants, thresholds, and non-obvious transformation functions — explain what the value represents and why it was chosen.
|
||||
- Comment fallback/guard code paths with when they trigger and why they exist.
|
||||
- Ensure doc comments match actual semantics — distinguish mutates-in-place (`&mut self`) from returns-new-value.
|
||||
- Use explicit forward-looking language (`TODO`, `FIXME`) in comments to distinguish current behavior from planned changes.
|
||||
- Document the semantic meaning of both present and absent states for `Option<T>` fields.
|
||||
- Use precise domain terminology — avoid ambiguous abbreviations (e.g., "FIXED" vs "fixed-width") or incorrect terms (e.g., "fields" when meaning "fragments").
|
||||
|
||||
## lance-encoding
|
||||
|
||||
Performance-critical encoding/decoding paths have additional requirements:
|
||||
|
||||
- Hoist loop-invariant conditionals out of hot loops — branch once outside, then use separate loop bodies or monomorphized variants.
|
||||
- Pre-allocate single contiguous buffers. Default to `buf.resize(len, 0)` for safe initialization; reserve `Vec::with_capacity` + `unsafe { set_len() }` for measured hot paths only, with a `// SAFETY:` comment explaining why the buffer will be fully initialized before read (e.g., immediately followed by `read_exact`).
|
||||
- Use `spawn_cpu()` only at the async-to-CPU boundary (e.g., FSST, decompression, batch materialization) — never nest redundant `spawn_cpu()` calls.
|
||||
- Use `expect_next()` and similar utility methods instead of inlining `None`-checks with error returns.
|
||||
Symlink
+1
@@ -0,0 +1 @@
|
||||
AGENTS.md
|
||||
@@ -0,0 +1,41 @@
|
||||
# Contributing to Rust
|
||||
|
||||
To format and lint Rust code:
|
||||
|
||||
```bash
|
||||
cargo fmt --all
|
||||
cargo clippy --all-features --tests --benches
|
||||
```
|
||||
|
||||
## Core Format
|
||||
|
||||
The core format is implemented in Rust under the `rust` directory. Once you've setup Rust you can build the core format with:
|
||||
|
||||
```bash
|
||||
cargo build
|
||||
```
|
||||
|
||||
This builds the debug build. For the optimized release build:
|
||||
|
||||
```bash
|
||||
cargo build -r
|
||||
```
|
||||
|
||||
To run the Rust unit tests:
|
||||
|
||||
```bash
|
||||
cargo test
|
||||
```
|
||||
|
||||
If you're working on a performance related feature, benchmarks can be run via:
|
||||
|
||||
```bash
|
||||
cargo bench
|
||||
```
|
||||
|
||||
If you want detailed logging and full backtraces, set the following environment variables.
|
||||
More details can be found [here](../docs/src/guide/performance.md#logging).
|
||||
|
||||
```bash
|
||||
LANCE_LOG=info RUST_BACKTRACE=FULL <cargo-commands>
|
||||
```
|
||||
@@ -0,0 +1,2 @@
|
||||
# Lance Rust Workspace
|
||||
Where core rust code lance lives
|
||||
@@ -0,0 +1,31 @@
|
||||
[package]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Arrow scalar type with Ord, Hash, and Eq support"
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
rust-version.workspace = true
|
||||
readme = "README.md"
|
||||
|
||||
[dependencies]
|
||||
# Note: this is a core crate and we should aim to keep this dependency list
|
||||
# as minimal as possible.
|
||||
arrow-array = { workspace = true }
|
||||
arrow-buffer = { workspace = true }
|
||||
arrow-cast = { workspace = true }
|
||||
arrow-data = { workspace = true }
|
||||
arrow-row = { workspace = true }
|
||||
arrow-schema = { workspace = true }
|
||||
half = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
arrow-ord = { workspace = true }
|
||||
proptest = { workspace = true }
|
||||
rstest = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,57 @@
|
||||
# lance-arrow-scalar
|
||||
|
||||
A scalar type backed by Apache Arrow arrays with `Ord`, `Hash`, and `Eq` support.
|
||||
|
||||
## Overview
|
||||
|
||||
`ArrowScalar` wraps a single-element Arrow array and provides comparison and hashing operations by leveraging Apache Arrow's `OwnedRow` representation. This ensures:
|
||||
|
||||
- **Correct total ordering** for all Arrow types
|
||||
- **Proper NaN handling** for floating-point values
|
||||
- **Consistent null ordering**
|
||||
- **O(1) comparisons** via cached row bytes
|
||||
|
||||
## Features
|
||||
|
||||
- `Eq`, `Ord`, and `Hash` traits for Arrow scalar values
|
||||
- Support for all Arrow data types
|
||||
- Serde serialization/deserialization support
|
||||
- Zero-copy conversion from Arrow arrays
|
||||
|
||||
## Usage
|
||||
|
||||
Add to your `Cargo.toml`:
|
||||
|
||||
```toml
|
||||
[dependencies]
|
||||
lance-arrow-scalar = "57.0.0"
|
||||
```
|
||||
|
||||
Then use in your code:
|
||||
|
||||
```rust
|
||||
use lance_arrow_scalar::ArrowScalar;
|
||||
|
||||
// Create from primitive types
|
||||
let a = ArrowScalar::from(42i32);
|
||||
let b = ArrowScalar::from(100i32);
|
||||
assert!(a < b);
|
||||
|
||||
// Create from strings
|
||||
let s1 = ArrowScalar::from("hello");
|
||||
let s2 = ArrowScalar::from("world");
|
||||
assert!(s1 < s2);
|
||||
|
||||
// Use in collections
|
||||
use std::collections::HashMap;
|
||||
let mut map = HashMap::new();
|
||||
map.insert(ArrowScalar::from("key"), ArrowScalar::from(123));
|
||||
```
|
||||
|
||||
## Cross-Type Comparison
|
||||
|
||||
Comparing scalars of different data types produces an arbitrary but consistent ordering based on the underlying row bytes. This allows scalars to be used as keys in sorted collections regardless of type, though the ordering across types is not semantically meaningful.
|
||||
|
||||
## Implementation Details
|
||||
|
||||
Comparisons and hashing are delegated to [`arrow_row::OwnedRow`], which provides efficient byte-level operations. The row representation is cached at construction time, making all comparison and hashing operations O(1).
|
||||
@@ -0,0 +1,108 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::*;
|
||||
use half::f16;
|
||||
|
||||
use crate::ArrowScalar;
|
||||
|
||||
macro_rules! impl_from_primitive {
|
||||
($native_ty:ty, $array_ty:ty) => {
|
||||
impl From<$native_ty> for ArrowScalar {
|
||||
fn from(value: $native_ty) -> Self {
|
||||
let array: ArrayRef = Arc::new(<$array_ty>::from(vec![value]));
|
||||
Self::try_from_array(array).expect("single-element primitive array is always valid")
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
impl_from_primitive!(i8, Int8Array);
|
||||
impl_from_primitive!(i16, Int16Array);
|
||||
impl_from_primitive!(i32, Int32Array);
|
||||
impl_from_primitive!(i64, Int64Array);
|
||||
impl_from_primitive!(u8, UInt8Array);
|
||||
impl_from_primitive!(u16, UInt16Array);
|
||||
impl_from_primitive!(u32, UInt32Array);
|
||||
impl_from_primitive!(u64, UInt64Array);
|
||||
impl_from_primitive!(f32, Float32Array);
|
||||
impl_from_primitive!(f64, Float64Array);
|
||||
|
||||
impl From<bool> for ArrowScalar {
|
||||
fn from(value: bool) -> Self {
|
||||
let array: ArrayRef = Arc::new(BooleanArray::from(vec![value]));
|
||||
Self::try_from_array(array).expect("single-element boolean array is always valid")
|
||||
}
|
||||
}
|
||||
|
||||
impl From<f16> for ArrowScalar {
|
||||
fn from(value: f16) -> Self {
|
||||
let array: ArrayRef = Arc::new(Float16Array::from(vec![value]));
|
||||
Self::try_from_array(array).expect("single-element f16 array is always valid")
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for ArrowScalar {
|
||||
fn from(value: &str) -> Self {
|
||||
let array: ArrayRef = Arc::new(StringArray::from(vec![value]));
|
||||
Self::try_from_array(array).expect("single-element string array is always valid")
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for ArrowScalar {
|
||||
fn from(value: String) -> Self {
|
||||
Self::from(value.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&[u8]> for ArrowScalar {
|
||||
fn from(value: &[u8]) -> Self {
|
||||
let array: ArrayRef = Arc::new(BinaryArray::from_vec(vec![value]));
|
||||
Self::try_from_array(array).expect("single-element binary array is always valid")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_from_primitives() {
|
||||
let s = ArrowScalar::from(42i32);
|
||||
assert!(!s.is_null());
|
||||
assert_eq!(format!("{s}"), "42");
|
||||
|
||||
let s = ArrowScalar::from(1.5f64);
|
||||
assert!(!s.is_null());
|
||||
|
||||
let s = ArrowScalar::from(true);
|
||||
assert_eq!(format!("{s}"), "true");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_string_types() {
|
||||
let s = ArrowScalar::from("hello");
|
||||
assert_eq!(format!("{s}"), "hello");
|
||||
|
||||
let s = ArrowScalar::from(String::from("world"));
|
||||
assert_eq!(format!("{s}"), "world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_binary() {
|
||||
let bytes: &[u8] = &[0xDE, 0xAD];
|
||||
let s = ArrowScalar::from(bytes);
|
||||
assert!(!s.is_null());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_f16() {
|
||||
let s = ArrowScalar::from(f16::from_f32(1.5));
|
||||
assert!(!s.is_null());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,624 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! A scalar type backed by a single-element Arrow array with [`Ord`], [`Hash`],
|
||||
//! and [`Eq`] support.
|
||||
//!
|
||||
//! Comparisons and hashing are delegated to [`arrow_row::OwnedRow`], which
|
||||
//! provides a correct total ordering for all Arrow types (including proper NaN
|
||||
//! handling for floats and null ordering).
|
||||
|
||||
mod convert;
|
||||
pub mod serde;
|
||||
|
||||
use std::cmp::Ordering;
|
||||
use std::fmt;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::cast::AsArray;
|
||||
use arrow_array::types::{Float16Type, Float32Type, Float64Type};
|
||||
use arrow_array::{ArrayRef, make_array, new_null_array};
|
||||
use arrow_cast::display::ArrayFormatter;
|
||||
use arrow_data::transform::MutableArrayData;
|
||||
use arrow_row::{OwnedRow, RowConverter, SortField};
|
||||
use arrow_schema::{ArrowError, DataType};
|
||||
|
||||
type Result<T> = std::result::Result<T, ArrowError>;
|
||||
|
||||
/// A scalar value backed by a length-1 Arrow array.
|
||||
///
|
||||
/// `ArrowScalar` provides [`Eq`], [`Ord`], and [`Hash`] by caching an
|
||||
/// [`OwnedRow`] at construction time. This means comparisons and hashing are
|
||||
/// O(1) row-byte operations rather than per-type dispatch.
|
||||
///
|
||||
/// # Cross-type comparison
|
||||
///
|
||||
/// Comparing scalars of different data types produces an arbitrary but
|
||||
/// consistent ordering based on the underlying row bytes. This is intentional
|
||||
/// — it allows scalars to be used as keys in sorted collections regardless of
|
||||
/// type, but the ordering across types is not semantically meaningful.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use lance_arrow_scalar::ArrowScalar;
|
||||
///
|
||||
/// let a = ArrowScalar::from(1i32);
|
||||
/// let b = ArrowScalar::from(2i32);
|
||||
/// assert!(a < b);
|
||||
///
|
||||
/// let c = ArrowScalar::from("hello");
|
||||
/// assert_eq!(c, ArrowScalar::from("hello"));
|
||||
/// ```
|
||||
pub struct ArrowScalar {
|
||||
array: ArrayRef,
|
||||
row: OwnedRow,
|
||||
}
|
||||
|
||||
impl ArrowScalar {
|
||||
/// Create a scalar by extracting the element at `offset` from `array`.
|
||||
pub fn try_new(array: &ArrayRef, offset: usize) -> Result<Self> {
|
||||
if offset >= array.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Scalar index out of bounds".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let data = array.to_data();
|
||||
let mut mutable = MutableArrayData::new(vec![&data], true, 1);
|
||||
mutable.extend(0, offset, offset + 1);
|
||||
let single = make_array(mutable.freeze());
|
||||
Self::try_from_array(single)
|
||||
}
|
||||
|
||||
/// Create a scalar from a length-1 array.
|
||||
pub fn try_from_array(array: ArrayRef) -> Result<Self> {
|
||||
if array.len() != 1 {
|
||||
return Err(ArrowError::InvalidArgumentError(format!(
|
||||
"ArrowScalar requires a length-1 array, got length {}",
|
||||
array.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let row = Self::compute_row(&array)?;
|
||||
Ok(Self { array, row })
|
||||
}
|
||||
|
||||
/// Create a null scalar of the given data type.
|
||||
pub fn new_null(data_type: &DataType) -> Result<Self> {
|
||||
Self::try_from_array(new_null_array(data_type, 1))
|
||||
}
|
||||
|
||||
fn compute_row(array: &ArrayRef) -> Result<OwnedRow> {
|
||||
let sort_field = SortField::new(array.data_type().clone());
|
||||
let converter = RowConverter::new(vec![sort_field])?;
|
||||
let rows = converter.convert_columns(&[Arc::clone(array)])?;
|
||||
Ok(rows.row(0).owned())
|
||||
}
|
||||
|
||||
/// Returns a reference to the underlying length-1 array.
|
||||
pub fn as_array(&self) -> &ArrayRef {
|
||||
&self.array
|
||||
}
|
||||
|
||||
/// Returns the data type of this scalar.
|
||||
pub fn data_type(&self) -> &DataType {
|
||||
self.array.data_type()
|
||||
}
|
||||
|
||||
/// Returns `true` if this scalar is null.
|
||||
pub fn is_null(&self) -> bool {
|
||||
self.array.null_count() == 1
|
||||
}
|
||||
|
||||
/// Returns `true` if this scalar is a non-null floating-point NaN.
|
||||
///
|
||||
/// ```
|
||||
/// use lance_arrow_scalar::ArrowScalar;
|
||||
///
|
||||
/// assert!(ArrowScalar::from(f32::NAN).is_nan());
|
||||
/// assert!(!ArrowScalar::from(1.0f32).is_nan());
|
||||
/// assert!(!ArrowScalar::from(1i32).is_nan());
|
||||
/// ```
|
||||
pub fn is_nan(&self) -> bool {
|
||||
if self.is_null() {
|
||||
return false;
|
||||
}
|
||||
|
||||
match self.data_type() {
|
||||
DataType::Float16 => self.array.as_primitive::<Float16Type>().value(0).is_nan(),
|
||||
DataType::Float32 => self.array.as_primitive::<Float32Type>().value(0).is_nan(),
|
||||
DataType::Float64 => self.array.as_primitive::<Float64Type>().value(0).is_nan(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for ArrowScalar {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.row == other.row
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for ArrowScalar {}
|
||||
|
||||
impl PartialOrd for ArrowScalar {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
|
||||
impl Ord for ArrowScalar {
|
||||
fn cmp(&self, other: &Self) -> Ordering {
|
||||
self.row.cmp(&other.row)
|
||||
}
|
||||
}
|
||||
|
||||
impl Hash for ArrowScalar {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
self.row.hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ArrowScalar {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
if self.is_null() {
|
||||
return write!(f, "null");
|
||||
}
|
||||
let formatter =
|
||||
ArrayFormatter::try_new(&self.array, &Default::default()).map_err(|_| fmt::Error)?;
|
||||
write!(f, "{}", formatter.value(0))
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for ArrowScalar {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "ArrowScalar({}: {})", self.data_type(), self)
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for ArrowScalar {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
array: Arc::clone(&self.array),
|
||||
row: self.row.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::{BTreeSet, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::*;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_try_new_extracts_element() {
|
||||
let array: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30]));
|
||||
let s = ArrowScalar::try_new(&array, 1).unwrap();
|
||||
assert_eq!(format!("{s}"), "20");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_try_new_out_of_bounds() {
|
||||
let array: ArrayRef = Arc::new(Int32Array::from(vec![1]));
|
||||
assert!(ArrowScalar::try_new(&array, 5).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_try_from_array_wrong_length() {
|
||||
let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
|
||||
assert!(ArrowScalar::try_from_array(array).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_equality() {
|
||||
let a = ArrowScalar::from(42i32);
|
||||
let b = ArrowScalar::from(42i32);
|
||||
let c = ArrowScalar::from(99i32);
|
||||
assert_eq!(a, b);
|
||||
assert_ne!(a, c);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ordering() {
|
||||
let a = ArrowScalar::from(1i32);
|
||||
let b = ArrowScalar::from(2i32);
|
||||
let c = ArrowScalar::from(3i32);
|
||||
assert!(a < b);
|
||||
assert!(b < c);
|
||||
assert_eq!(a.cmp(&a), Ordering::Equal);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hash_consistent_with_eq() {
|
||||
use std::hash::DefaultHasher;
|
||||
|
||||
let a = ArrowScalar::from(42i32);
|
||||
let b = ArrowScalar::from(42i32);
|
||||
let hash_a = {
|
||||
let mut h = DefaultHasher::new();
|
||||
a.hash(&mut h);
|
||||
h.finish()
|
||||
};
|
||||
let hash_b = {
|
||||
let mut h = DefaultHasher::new();
|
||||
b.hash(&mut h);
|
||||
h.finish()
|
||||
};
|
||||
assert_eq!(hash_a, hash_b);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_in_hashset() {
|
||||
let mut set = HashSet::new();
|
||||
set.insert(ArrowScalar::from(1i32));
|
||||
set.insert(ArrowScalar::from(2i32));
|
||||
set.insert(ArrowScalar::from(1i32));
|
||||
assert_eq!(set.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_in_btreeset() {
|
||||
let mut set = BTreeSet::new();
|
||||
set.insert(ArrowScalar::from(3i32));
|
||||
set.insert(ArrowScalar::from(1i32));
|
||||
set.insert(ArrowScalar::from(2i32));
|
||||
let values: Vec<_> = set.iter().map(|s| format!("{s}")).collect();
|
||||
assert_eq!(values, vec!["1", "2", "3"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_null_scalar() {
|
||||
let array: ArrayRef = Arc::new(Int32Array::from(vec![None]));
|
||||
let s = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert!(s.is_null());
|
||||
assert_eq!(format!("{s}"), "null");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_null_sorts_first() {
|
||||
let null_scalar = {
|
||||
let array: ArrayRef = Arc::new(Int32Array::from(vec![None]));
|
||||
ArrowScalar::try_from_array(array).unwrap()
|
||||
};
|
||||
let value_scalar = ArrowScalar::from(0i32);
|
||||
assert!(null_scalar < value_scalar);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::float_nan(
|
||||
ArrowScalar::from(f64::NAN),
|
||||
ArrowScalar::from(f64::INFINITY),
|
||||
Ordering::Greater
|
||||
)]
|
||||
#[case::float_normal(ArrowScalar::from(1.0f64), ArrowScalar::from(2.0f64), Ordering::Less)]
|
||||
fn test_float_ordering(
|
||||
#[case] a: ArrowScalar,
|
||||
#[case] b: ArrowScalar,
|
||||
#[case] expected: Ordering,
|
||||
) {
|
||||
assert_eq!(a.cmp(&b), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::float16_nan(ArrowScalar::from(half::f16::NAN), true)]
|
||||
#[case::float32_nan(ArrowScalar::from(f32::NAN), true)]
|
||||
#[case::float64_nan(ArrowScalar::from(f64::NAN), true)]
|
||||
#[case::float64_finite(ArrowScalar::from(1.0f64), false)]
|
||||
#[case::int32(ArrowScalar::from(1i32), false)]
|
||||
fn test_is_nan(#[case] scalar: ArrowScalar, #[case] expected: bool) {
|
||||
assert_eq!(scalar.is_nan(), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_null_is_not_nan() {
|
||||
let array: ArrayRef = Arc::new(Float64Array::from(vec![None]));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert!(!scalar.is_nan());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_display_string() {
|
||||
let s = ArrowScalar::from("hello world");
|
||||
assert_eq!(format!("{s}"), "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_debug() {
|
||||
let s = ArrowScalar::from(42i32);
|
||||
let debug = format!("{s:?}");
|
||||
assert!(debug.contains("ArrowScalar"));
|
||||
assert!(debug.contains("42"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clone() {
|
||||
let a = ArrowScalar::from(42i32);
|
||||
let b = a.clone();
|
||||
assert_eq!(a, b);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_data_type() {
|
||||
let s = ArrowScalar::from(42i32);
|
||||
assert_eq!(s.data_type(), &DataType::Int32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_boolean_roundtrip() {
|
||||
let t = ArrowScalar::from(true);
|
||||
let f = ArrowScalar::from(false);
|
||||
assert_eq!(t.data_type(), &DataType::Boolean);
|
||||
assert!(!t.is_null());
|
||||
assert_eq!(format!("{t}"), "true");
|
||||
assert_eq!(format!("{f}"), "false");
|
||||
|
||||
// Extract from multi-element array
|
||||
let array: ArrayRef = Arc::new(BooleanArray::from(vec![true, false, true]));
|
||||
let s = ArrowScalar::try_new(&array, 1).unwrap();
|
||||
assert_eq!(format!("{s}"), "false");
|
||||
assert_eq!(s.data_type(), &DataType::Boolean);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_boolean_equality_and_ordering() {
|
||||
let t1 = ArrowScalar::from(true);
|
||||
let t2 = ArrowScalar::from(true);
|
||||
let f1 = ArrowScalar::from(false);
|
||||
assert_eq!(t1, t2);
|
||||
assert_ne!(t1, f1);
|
||||
// false < true in arrow row encoding
|
||||
assert!(f1 < t1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_boolean_null() {
|
||||
let array: ArrayRef = Arc::new(BooleanArray::from(vec![None]));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert!(scalar.is_null());
|
||||
assert_eq!(scalar.data_type(), &DataType::Boolean);
|
||||
assert_eq!(format!("{scalar}"), "null");
|
||||
|
||||
// null sorts before false
|
||||
let f = ArrowScalar::from(false);
|
||||
assert!(scalar < f);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_string_view_roundtrip() {
|
||||
let array: ArrayRef = Arc::new(StringViewArray::from(vec![
|
||||
"hello world, this is a long string view",
|
||||
]));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert_eq!(scalar.data_type(), &DataType::Utf8View);
|
||||
assert!(!scalar.is_null());
|
||||
assert_eq!(
|
||||
format!("{scalar}"),
|
||||
"hello world, this is a long string view"
|
||||
);
|
||||
|
||||
// Extract from multi-element array
|
||||
let array: ArrayRef = Arc::new(StringViewArray::from(vec!["alpha", "beta", "gamma"]));
|
||||
let s = ArrowScalar::try_new(&array, 1).unwrap();
|
||||
assert_eq!(format!("{s}"), "beta");
|
||||
assert_eq!(s.data_type(), &DataType::Utf8View);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_binary_view_roundtrip() {
|
||||
let values: Vec<&[u8]> = vec![b"\xDE\xAD\xBE\xEF"];
|
||||
let array: ArrayRef = Arc::new(BinaryViewArray::from(values));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert_eq!(scalar.data_type(), &DataType::BinaryView);
|
||||
assert!(!scalar.is_null());
|
||||
|
||||
// Extract from multi-element array
|
||||
let values: Vec<&[u8]> = vec![b"aaa", b"bbb", b"ccc"];
|
||||
let array: ArrayRef = Arc::new(BinaryViewArray::from(values));
|
||||
let s = ArrowScalar::try_new(&array, 2).unwrap();
|
||||
assert_eq!(s.data_type(), &DataType::BinaryView);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_string_view_equality_and_ordering() {
|
||||
let mk = |s: &str| {
|
||||
let array: ArrayRef = Arc::new(StringViewArray::from(vec![s]));
|
||||
ArrowScalar::try_from_array(array).unwrap()
|
||||
};
|
||||
let a = mk("apple");
|
||||
let b = mk("apple");
|
||||
let c = mk("banana");
|
||||
assert_eq!(a, b);
|
||||
assert_ne!(a, c);
|
||||
assert!(a < c);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_binary_view_equality_and_ordering() {
|
||||
let mk = |b: &[u8]| {
|
||||
let values: Vec<&[u8]> = vec![b];
|
||||
let array: ArrayRef = Arc::new(BinaryViewArray::from(values));
|
||||
ArrowScalar::try_from_array(array).unwrap()
|
||||
};
|
||||
let a = mk(b"\x01\x02");
|
||||
let b = mk(b"\x01\x02");
|
||||
let c = mk(b"\x01\x03");
|
||||
assert_eq!(a, b);
|
||||
assert_ne!(a, c);
|
||||
assert!(a < c);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_string_view_in_collections() {
|
||||
let mk = |s: &str| {
|
||||
let array: ArrayRef = Arc::new(StringViewArray::from(vec![s]));
|
||||
ArrowScalar::try_from_array(array).unwrap()
|
||||
};
|
||||
|
||||
let mut hset = HashSet::new();
|
||||
hset.insert(mk("foo"));
|
||||
hset.insert(mk("bar"));
|
||||
hset.insert(mk("foo"));
|
||||
assert_eq!(hset.len(), 2);
|
||||
|
||||
let mut bset = BTreeSet::new();
|
||||
bset.insert(mk("cherry"));
|
||||
bset.insert(mk("apple"));
|
||||
bset.insert(mk("banana"));
|
||||
let sorted: Vec<_> = bset.iter().map(|s| format!("{s}")).collect();
|
||||
assert_eq!(sorted, vec!["apple", "banana", "cherry"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_string_view_null() {
|
||||
let array: ArrayRef = Arc::new(StringViewArray::from(vec![Option::<&str>::None]));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert!(scalar.is_null());
|
||||
assert_eq!(scalar.data_type(), &DataType::Utf8View);
|
||||
assert_eq!(format!("{scalar}"), "null");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_binary_view_null() {
|
||||
let array: ArrayRef = Arc::new(BinaryViewArray::from(vec![Option::<&[u8]>::None]));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert!(scalar.is_null());
|
||||
assert_eq!(scalar.data_type(), &DataType::BinaryView);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cross_type_comparison_is_consistent() {
|
||||
let int_scalar = ArrowScalar::from(42i32);
|
||||
let str_scalar = ArrowScalar::from("hello");
|
||||
// The ordering is arbitrary but must be consistent
|
||||
let ord1 = int_scalar.cmp(&str_scalar);
|
||||
let ord2 = int_scalar.cmp(&str_scalar);
|
||||
assert_eq!(ord1, ord2);
|
||||
// And the reverse should be opposite
|
||||
assert_eq!(str_scalar.cmp(&int_scalar), ord1.reverse());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod prop_tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::*;
|
||||
use arrow_ord::sort::sort;
|
||||
use arrow_schema::SortOptions;
|
||||
use proptest::prelude::*;
|
||||
|
||||
use super::ArrowScalar;
|
||||
|
||||
/// Generate an arbitrary Arrow array of a randomly chosen type, including
|
||||
/// nulls. Covers primitives, booleans, string/binary types and their view
|
||||
/// variants.
|
||||
fn arbitrary_array() -> BoxedStrategy<ArrayRef> {
|
||||
let len = 0..=100usize;
|
||||
|
||||
prop_oneof![
|
||||
// --- integer types ---
|
||||
proptest::collection::vec(proptest::option::of(any::<i8>()), len.clone())
|
||||
.prop_map(|v| Arc::new(Int8Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<i16>()), len.clone())
|
||||
.prop_map(|v| Arc::new(Int16Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<i32>()), len.clone())
|
||||
.prop_map(|v| Arc::new(Int32Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<i64>()), len.clone())
|
||||
.prop_map(|v| Arc::new(Int64Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<u8>()), len.clone())
|
||||
.prop_map(|v| Arc::new(UInt8Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<u16>()), len.clone())
|
||||
.prop_map(|v| Arc::new(UInt16Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<u32>()), len.clone())
|
||||
.prop_map(|v| Arc::new(UInt32Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<u64>()), len.clone())
|
||||
.prop_map(|v| Arc::new(UInt64Array::from(v)) as ArrayRef),
|
||||
// --- float types ---
|
||||
proptest::collection::vec(proptest::option::of(any::<f32>()), len.clone())
|
||||
.prop_map(|v| Arc::new(Float32Array::from(v)) as ArrayRef),
|
||||
proptest::collection::vec(proptest::option::of(any::<f64>()), len.clone())
|
||||
.prop_map(|v| Arc::new(Float64Array::from(v)) as ArrayRef),
|
||||
// --- boolean ---
|
||||
proptest::collection::vec(proptest::option::of(any::<bool>()), len.clone())
|
||||
.prop_map(|v| Arc::new(BooleanArray::from(v)) as ArrayRef),
|
||||
// --- string types ---
|
||||
proptest::collection::vec(proptest::option::of(any::<String>()), len.clone()).prop_map(
|
||||
|v| {
|
||||
let refs: Vec<Option<&str>> = v.iter().map(|o| o.as_deref()).collect();
|
||||
Arc::new(StringArray::from(refs)) as ArrayRef
|
||||
}
|
||||
),
|
||||
proptest::collection::vec(proptest::option::of(any::<String>()), len.clone()).prop_map(
|
||||
|v| {
|
||||
let refs: Vec<Option<&str>> = v.iter().map(|o| o.as_deref()).collect();
|
||||
Arc::new(LargeStringArray::from(refs)) as ArrayRef
|
||||
}
|
||||
),
|
||||
proptest::collection::vec(proptest::option::of(any::<String>()), len.clone()).prop_map(
|
||||
|v| {
|
||||
let refs: Vec<Option<&str>> = v.iter().map(|o| o.as_deref()).collect();
|
||||
Arc::new(StringViewArray::from(refs)) as ArrayRef
|
||||
}
|
||||
),
|
||||
// --- binary types ---
|
||||
proptest::collection::vec(
|
||||
proptest::option::of(proptest::collection::vec(any::<u8>(), 0..50)),
|
||||
len.clone(),
|
||||
)
|
||||
.prop_map(|v| {
|
||||
let refs: Vec<Option<&[u8]>> = v.iter().map(|o| o.as_deref()).collect();
|
||||
Arc::new(BinaryArray::from(refs)) as ArrayRef
|
||||
}),
|
||||
proptest::collection::vec(
|
||||
proptest::option::of(proptest::collection::vec(any::<u8>(), 0..50)),
|
||||
len.clone(),
|
||||
)
|
||||
.prop_map(|v| {
|
||||
let refs: Vec<Option<&[u8]>> = v.iter().map(|o| o.as_deref()).collect();
|
||||
Arc::new(LargeBinaryArray::from(refs)) as ArrayRef
|
||||
}),
|
||||
proptest::collection::vec(
|
||||
proptest::option::of(proptest::collection::vec(any::<u8>(), 0..50)),
|
||||
len,
|
||||
)
|
||||
.prop_map(|v| {
|
||||
let refs: Vec<Option<&[u8]>> = v.iter().map(|o| o.as_deref()).collect();
|
||||
Arc::new(BinaryViewArray::from(refs)) as ArrayRef
|
||||
}),
|
||||
]
|
||||
.boxed()
|
||||
}
|
||||
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn sorted_array_produces_sorted_scalars(array in arbitrary_array()) {
|
||||
let sorted = sort(
|
||||
&array,
|
||||
Some(SortOptions { descending: false, nulls_first: true }),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let scalars: Vec<ArrowScalar> = (0..sorted.len())
|
||||
.map(|i| ArrowScalar::try_new(&sorted, i).unwrap())
|
||||
.collect();
|
||||
|
||||
for i in 1..scalars.len() {
|
||||
prop_assert!(
|
||||
scalars[i - 1] <= scalars[i],
|
||||
"scalar[{}] ({:?}) should be <= scalar[{}] ({:?})",
|
||||
i - 1, scalars[i - 1], i, scalars[i],
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,561 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Binary serialization for [`ArrowScalar`].
|
||||
//!
|
||||
//! Default format (with type prefix):
|
||||
//! ```text
|
||||
//! | varint: format_string_len | raw: format_string_bytes |
|
||||
//! | varint: null_flag (0 = non-null, 1 = null) |
|
||||
//! | varint: num_buffers | (only if non-null)
|
||||
//! | varint: buffer_0_len | ... | varint: buffer_{n-1}_len | (only if non-null)
|
||||
//! | raw: buffer_0 bytes | ... | raw: buffer_{n-1} bytes | (only if non-null)
|
||||
//! ```
|
||||
//!
|
||||
//! The format string uses the
|
||||
//! [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html#data-type-description-format-strings)
|
||||
//! encoding. Use [`EncodeOptions`] / [`DecodeOptions`] to omit the type prefix
|
||||
//! when the caller already knows the data type.
|
||||
|
||||
use std::borrow::Cow;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::make_array;
|
||||
use arrow_buffer::Buffer;
|
||||
use arrow_data::ArrayDataBuilder;
|
||||
use arrow_schema::{ArrowError, DataType, IntervalUnit, TimeUnit};
|
||||
|
||||
use crate::ArrowScalar;
|
||||
|
||||
type Result<T> = std::result::Result<T, ArrowError>;
|
||||
|
||||
/// Options for [`ArrowScalar::encode_with_options`].
|
||||
pub struct EncodeOptions {
|
||||
/// When `true` (the default), the Arrow C Data Interface format string
|
||||
/// for the scalar's data type is prepended as a varint-length-prefixed
|
||||
/// UTF-8 string. Set to `false` to omit the type prefix (the caller
|
||||
/// must then supply the `DataType` at decode time).
|
||||
pub include_data_type: bool,
|
||||
}
|
||||
|
||||
impl Default for EncodeOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
include_data_type: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Options for [`ArrowScalar::decode_with_options`].
|
||||
#[derive(Default)]
|
||||
pub struct DecodeOptions<'a> {
|
||||
/// When `Some`, the data type is taken from this value and the encoded
|
||||
/// bytes are assumed to contain no type prefix. When `None` (the
|
||||
/// default), the data type is read from the encoded format-string prefix.
|
||||
pub data_type: Option<&'a DataType>,
|
||||
}
|
||||
|
||||
/// Encode a `u64` as a variable-length integer (LEB128).
|
||||
///
|
||||
/// Values below 128 use a single byte; the maximum encoding is 10 bytes.
|
||||
pub fn encode_varint(out: &mut Vec<u8>, mut value: u64) {
|
||||
loop {
|
||||
let byte = (value & 0x7F) as u8;
|
||||
value >>= 7;
|
||||
if value == 0 {
|
||||
out.push(byte);
|
||||
return;
|
||||
}
|
||||
out.push(byte | 0x80);
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode a variable-length integer (LEB128) from `buf` at the given `offset`.
|
||||
///
|
||||
/// On success, `offset` is advanced past the consumed bytes.
|
||||
pub fn decode_varint(buf: &[u8], offset: &mut usize) -> Result<u64> {
|
||||
let mut result: u64 = 0;
|
||||
let mut shift = 0u32;
|
||||
loop {
|
||||
if *offset >= buf.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid varint: unexpected EOF".to_string(),
|
||||
));
|
||||
}
|
||||
let byte = buf[*offset];
|
||||
*offset += 1;
|
||||
|
||||
result |= u64::from(byte & 0x7F) << shift;
|
||||
if byte & 0x80 == 0 {
|
||||
return Ok(result);
|
||||
}
|
||||
shift += 7;
|
||||
if shift >= 64 {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid varint: too many bytes".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a [`DataType`] to its Arrow C Data Interface format string.
|
||||
///
|
||||
/// Only non-nested types are supported (nested types are already rejected by
|
||||
/// [`ArrowScalar::encode`]).
|
||||
fn data_type_to_format_string(dtype: &DataType) -> Result<Cow<'static, str>> {
|
||||
match dtype {
|
||||
DataType::Null => Ok("n".into()),
|
||||
DataType::Boolean => Ok("b".into()),
|
||||
DataType::Int8 => Ok("c".into()),
|
||||
DataType::UInt8 => Ok("C".into()),
|
||||
DataType::Int16 => Ok("s".into()),
|
||||
DataType::UInt16 => Ok("S".into()),
|
||||
DataType::Int32 => Ok("i".into()),
|
||||
DataType::UInt32 => Ok("I".into()),
|
||||
DataType::Int64 => Ok("l".into()),
|
||||
DataType::UInt64 => Ok("L".into()),
|
||||
DataType::Float16 => Ok("e".into()),
|
||||
DataType::Float32 => Ok("f".into()),
|
||||
DataType::Float64 => Ok("g".into()),
|
||||
DataType::Binary => Ok("z".into()),
|
||||
DataType::LargeBinary => Ok("Z".into()),
|
||||
DataType::Utf8 => Ok("u".into()),
|
||||
DataType::LargeUtf8 => Ok("U".into()),
|
||||
DataType::BinaryView => Ok("vz".into()),
|
||||
DataType::Utf8View => Ok("vu".into()),
|
||||
DataType::FixedSizeBinary(n) => Ok(Cow::Owned(format!("w:{n}"))),
|
||||
DataType::Decimal32(p, s) => Ok(Cow::Owned(format!("d:{p},{s},32"))),
|
||||
DataType::Decimal64(p, s) => Ok(Cow::Owned(format!("d:{p},{s},64"))),
|
||||
DataType::Decimal128(p, s) => Ok(Cow::Owned(format!("d:{p},{s}"))),
|
||||
DataType::Decimal256(p, s) => Ok(Cow::Owned(format!("d:{p},{s},256"))),
|
||||
DataType::Date32 => Ok("tdD".into()),
|
||||
DataType::Date64 => Ok("tdm".into()),
|
||||
DataType::Time32(TimeUnit::Second) => Ok("tts".into()),
|
||||
DataType::Time32(TimeUnit::Millisecond) => Ok("ttm".into()),
|
||||
DataType::Time64(TimeUnit::Microsecond) => Ok("ttu".into()),
|
||||
DataType::Time64(TimeUnit::Nanosecond) => Ok("ttn".into()),
|
||||
DataType::Timestamp(TimeUnit::Second, None) => Ok("tss:".into()),
|
||||
DataType::Timestamp(TimeUnit::Millisecond, None) => Ok("tsm:".into()),
|
||||
DataType::Timestamp(TimeUnit::Microsecond, None) => Ok("tsu:".into()),
|
||||
DataType::Timestamp(TimeUnit::Nanosecond, None) => Ok("tsn:".into()),
|
||||
DataType::Timestamp(TimeUnit::Second, Some(tz)) => Ok(Cow::Owned(format!("tss:{tz}"))),
|
||||
DataType::Timestamp(TimeUnit::Millisecond, Some(tz)) => Ok(Cow::Owned(format!("tsm:{tz}"))),
|
||||
DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => Ok(Cow::Owned(format!("tsu:{tz}"))),
|
||||
DataType::Timestamp(TimeUnit::Nanosecond, Some(tz)) => Ok(Cow::Owned(format!("tsn:{tz}"))),
|
||||
DataType::Duration(TimeUnit::Second) => Ok("tDs".into()),
|
||||
DataType::Duration(TimeUnit::Millisecond) => Ok("tDm".into()),
|
||||
DataType::Duration(TimeUnit::Microsecond) => Ok("tDu".into()),
|
||||
DataType::Duration(TimeUnit::Nanosecond) => Ok("tDn".into()),
|
||||
DataType::Interval(IntervalUnit::YearMonth) => Ok("tiM".into()),
|
||||
DataType::Interval(IntervalUnit::DayTime) => Ok("tiD".into()),
|
||||
DataType::Interval(IntervalUnit::MonthDayNano) => Ok("tin".into()),
|
||||
other => Err(ArrowError::InvalidArgumentError(format!(
|
||||
"Cannot encode data type as format string: {other:?}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse an Arrow C Data Interface format string back to a [`DataType`].
|
||||
///
|
||||
/// Only non-nested types are supported.
|
||||
fn format_string_to_data_type(fmt: &str) -> Result<DataType> {
|
||||
match fmt {
|
||||
"n" => Ok(DataType::Null),
|
||||
"b" => Ok(DataType::Boolean),
|
||||
"c" => Ok(DataType::Int8),
|
||||
"C" => Ok(DataType::UInt8),
|
||||
"s" => Ok(DataType::Int16),
|
||||
"S" => Ok(DataType::UInt16),
|
||||
"i" => Ok(DataType::Int32),
|
||||
"I" => Ok(DataType::UInt32),
|
||||
"l" => Ok(DataType::Int64),
|
||||
"L" => Ok(DataType::UInt64),
|
||||
"e" => Ok(DataType::Float16),
|
||||
"f" => Ok(DataType::Float32),
|
||||
"g" => Ok(DataType::Float64),
|
||||
"z" => Ok(DataType::Binary),
|
||||
"Z" => Ok(DataType::LargeBinary),
|
||||
"u" => Ok(DataType::Utf8),
|
||||
"U" => Ok(DataType::LargeUtf8),
|
||||
"vz" => Ok(DataType::BinaryView),
|
||||
"vu" => Ok(DataType::Utf8View),
|
||||
"tdD" => Ok(DataType::Date32),
|
||||
"tdm" => Ok(DataType::Date64),
|
||||
"tts" => Ok(DataType::Time32(TimeUnit::Second)),
|
||||
"ttm" => Ok(DataType::Time32(TimeUnit::Millisecond)),
|
||||
"ttu" => Ok(DataType::Time64(TimeUnit::Microsecond)),
|
||||
"ttn" => Ok(DataType::Time64(TimeUnit::Nanosecond)),
|
||||
"tDs" => Ok(DataType::Duration(TimeUnit::Second)),
|
||||
"tDm" => Ok(DataType::Duration(TimeUnit::Millisecond)),
|
||||
"tDu" => Ok(DataType::Duration(TimeUnit::Microsecond)),
|
||||
"tDn" => Ok(DataType::Duration(TimeUnit::Nanosecond)),
|
||||
"tiM" => Ok(DataType::Interval(IntervalUnit::YearMonth)),
|
||||
"tiD" => Ok(DataType::Interval(IntervalUnit::DayTime)),
|
||||
"tin" => Ok(DataType::Interval(IntervalUnit::MonthDayNano)),
|
||||
other => {
|
||||
let parts: Vec<&str> = other.splitn(2, ':').collect();
|
||||
match parts.as_slice() {
|
||||
["w", num_bytes] => {
|
||||
let n = num_bytes.parse::<i32>().map_err(|_| {
|
||||
ArrowError::InvalidArgumentError(
|
||||
"FixedSizeBinary requires an integer byte count".to_string(),
|
||||
)
|
||||
})?;
|
||||
Ok(DataType::FixedSizeBinary(n))
|
||||
}
|
||||
["d", extra] => {
|
||||
let dec_parts: Vec<&str> = extra.splitn(3, ',').collect();
|
||||
match dec_parts.as_slice() {
|
||||
[precision, scale] => {
|
||||
let p = precision.parse::<u8>().map_err(|_| {
|
||||
ArrowError::InvalidArgumentError(
|
||||
"Decimal requires an integer precision".to_string(),
|
||||
)
|
||||
})?;
|
||||
let s = scale.parse::<i8>().map_err(|_| {
|
||||
ArrowError::InvalidArgumentError(
|
||||
"Decimal requires an integer scale".to_string(),
|
||||
)
|
||||
})?;
|
||||
Ok(DataType::Decimal128(p, s))
|
||||
}
|
||||
[precision, scale, bits] => {
|
||||
let p = precision.parse::<u8>().map_err(|_| {
|
||||
ArrowError::InvalidArgumentError(
|
||||
"Decimal requires an integer precision".to_string(),
|
||||
)
|
||||
})?;
|
||||
let s = scale.parse::<i8>().map_err(|_| {
|
||||
ArrowError::InvalidArgumentError(
|
||||
"Decimal requires an integer scale".to_string(),
|
||||
)
|
||||
})?;
|
||||
match *bits {
|
||||
"32" => Ok(DataType::Decimal32(p, s)),
|
||||
"64" => Ok(DataType::Decimal64(p, s)),
|
||||
"128" => Ok(DataType::Decimal128(p, s)),
|
||||
"256" => Ok(DataType::Decimal256(p, s)),
|
||||
_ => Err(ArrowError::InvalidArgumentError(format!(
|
||||
"Unsupported decimal bit width: {bits}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
_ => Err(ArrowError::InvalidArgumentError(format!(
|
||||
"Invalid decimal format string: d:{extra}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
["tss", ""] => Ok(DataType::Timestamp(TimeUnit::Second, None)),
|
||||
["tsm", ""] => Ok(DataType::Timestamp(TimeUnit::Millisecond, None)),
|
||||
["tsu", ""] => Ok(DataType::Timestamp(TimeUnit::Microsecond, None)),
|
||||
["tsn", ""] => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)),
|
||||
["tss", tz] => Ok(DataType::Timestamp(TimeUnit::Second, Some(Arc::from(*tz)))),
|
||||
["tsm", tz] => Ok(DataType::Timestamp(
|
||||
TimeUnit::Millisecond,
|
||||
Some(Arc::from(*tz)),
|
||||
)),
|
||||
["tsu", tz] => Ok(DataType::Timestamp(
|
||||
TimeUnit::Microsecond,
|
||||
Some(Arc::from(*tz)),
|
||||
)),
|
||||
["tsn", tz] => Ok(DataType::Timestamp(
|
||||
TimeUnit::Nanosecond,
|
||||
Some(Arc::from(*tz)),
|
||||
)),
|
||||
_ => Err(ArrowError::InvalidArgumentError(format!(
|
||||
"Unsupported format string: {other:?}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ArrowScalar {
|
||||
/// Serialize this scalar to a self-describing binary representation.
|
||||
///
|
||||
/// The data type is encoded as a format-string prefix so that
|
||||
/// [`decode`](Self::decode) can reconstruct the scalar without external
|
||||
/// type information. Use [`encode_with_options`](Self::encode_with_options)
|
||||
/// to omit the prefix when the caller already knows the type.
|
||||
///
|
||||
/// Only non-nested scalars are supported. Null scalars are encoded as a
|
||||
/// null flag with no buffer data.
|
||||
pub fn encode(&self) -> Result<Vec<u8>> {
|
||||
self.encode_with_options(&EncodeOptions::default())
|
||||
}
|
||||
|
||||
/// Serialize this scalar with the given [`EncodeOptions`].
|
||||
pub fn encode_with_options(&self, options: &EncodeOptions) -> Result<Vec<u8>> {
|
||||
let array = self.as_array();
|
||||
let data = array.to_data();
|
||||
if !data.child_data().is_empty() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Cannot encode nested scalar".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut out = Vec::with_capacity(64);
|
||||
|
||||
if options.include_data_type {
|
||||
let fmt = data_type_to_format_string(array.data_type())?;
|
||||
encode_varint(&mut out, fmt.len() as u64);
|
||||
out.extend_from_slice(fmt.as_bytes());
|
||||
}
|
||||
|
||||
if self.is_null() {
|
||||
encode_varint(&mut out, 1); // null_flag = 1
|
||||
} else {
|
||||
encode_varint(&mut out, 0); // null_flag = 0
|
||||
let buffers = data.buffers();
|
||||
encode_varint(&mut out, buffers.len() as u64);
|
||||
for b in buffers {
|
||||
encode_varint(&mut out, b.len() as u64);
|
||||
}
|
||||
for b in buffers {
|
||||
out.extend_from_slice(b.as_slice());
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Deserialize a scalar from the self-describing binary representation
|
||||
/// produced by [`encode`](Self::encode).
|
||||
///
|
||||
/// The data type is read from the format-string prefix in the encoded
|
||||
/// bytes. Use [`decode_with_options`](Self::decode_with_options) to supply
|
||||
/// the type externally when the prefix was omitted at encode time.
|
||||
pub fn decode(buf: &[u8]) -> Result<Self> {
|
||||
Self::decode_with_options(buf, &DecodeOptions::default())
|
||||
}
|
||||
|
||||
/// Deserialize a scalar with the given [`DecodeOptions`].
|
||||
pub fn decode_with_options(buf: &[u8], options: &DecodeOptions) -> Result<Self> {
|
||||
let mut offset = 0;
|
||||
|
||||
let data_type = match options.data_type {
|
||||
Some(dt) => dt.clone(),
|
||||
None => {
|
||||
let fmt_len = decode_varint(buf, &mut offset)? as usize;
|
||||
if offset + fmt_len > buf.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid scalar buffer: unexpected EOF reading format string".to_string(),
|
||||
));
|
||||
}
|
||||
let fmt_str = std::str::from_utf8(&buf[offset..offset + fmt_len]).map_err(|e| {
|
||||
ArrowError::InvalidArgumentError(format!(
|
||||
"Invalid format string: not valid UTF-8: {e}"
|
||||
))
|
||||
})?;
|
||||
offset += fmt_len;
|
||||
format_string_to_data_type(fmt_str)?
|
||||
}
|
||||
};
|
||||
|
||||
let null_flag = decode_varint(buf, &mut offset)?;
|
||||
if null_flag == 1 {
|
||||
if offset != buf.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid scalar buffer: trailing bytes after null flag".to_string(),
|
||||
));
|
||||
}
|
||||
return Self::new_null(&data_type);
|
||||
}
|
||||
|
||||
let num_buffers = decode_varint(buf, &mut offset)? as usize;
|
||||
|
||||
let mut buffer_lens = Vec::with_capacity(num_buffers);
|
||||
for _ in 0..num_buffers {
|
||||
buffer_lens.push(decode_varint(buf, &mut offset)? as usize);
|
||||
}
|
||||
|
||||
let mut buffers = Vec::with_capacity(num_buffers);
|
||||
for len in &buffer_lens {
|
||||
if offset + len > buf.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid scalar buffer: unexpected EOF".to_string(),
|
||||
));
|
||||
}
|
||||
buffers.push(Buffer::from_vec(buf[offset..offset + len].to_vec()));
|
||||
offset += len;
|
||||
}
|
||||
|
||||
if offset != buf.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid scalar buffer: trailing bytes".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut builder = ArrayDataBuilder::new(data_type).len(1).null_count(0);
|
||||
for b in buffers {
|
||||
builder = builder.add_buffer(b);
|
||||
}
|
||||
let array = make_array(builder.build()?);
|
||||
Self::try_from_array(array)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{
|
||||
ArrayRef, BinaryViewArray, Int32Array, StringArray, StringViewArray,
|
||||
TimestampMicrosecondArray,
|
||||
};
|
||||
use arrow_schema::DataType;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::ArrowScalar;
|
||||
|
||||
#[test]
|
||||
fn test_varint_roundtrip() {
|
||||
for value in [0u64, 1, 127, 128, 16383, 16384, u64::MAX] {
|
||||
let mut buf = Vec::new();
|
||||
encode_varint(&mut buf, value);
|
||||
let mut offset = 0;
|
||||
let decoded = decode_varint(&buf, &mut offset).unwrap();
|
||||
assert_eq!(decoded, value);
|
||||
assert_eq!(offset, buf.len());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_varint_small_is_one_byte() {
|
||||
let mut buf = Vec::new();
|
||||
encode_varint(&mut buf, 42);
|
||||
assert_eq!(buf.len(), 1);
|
||||
assert_eq!(buf[0], 42);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::int32(Arc::new(Int32Array::from(vec![42])) as ArrayRef)]
|
||||
#[case::string(Arc::new(StringArray::from(vec!["hello"])) as ArrayRef)]
|
||||
#[case::string_view(Arc::new(StringViewArray::from(vec!["hello world, long string view"])) as ArrayRef)]
|
||||
#[case::binary_view(Arc::new(BinaryViewArray::from(vec![b"\xDE\xAD\xBE\xEF".as_ref()])) as ArrayRef)]
|
||||
fn test_encode_decode_roundtrip(#[case] array: ArrayRef) {
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
let encoded = scalar.encode().unwrap();
|
||||
let decoded = ArrowScalar::decode(&encoded).unwrap();
|
||||
assert_eq!(scalar, decoded);
|
||||
assert_eq!(scalar.data_type(), decoded.data_type());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::int32(Arc::new(Int32Array::from(vec![42])) as ArrayRef, DataType::Int32)]
|
||||
#[case::string(Arc::new(StringArray::from(vec!["hello"])) as ArrayRef, DataType::Utf8)]
|
||||
#[case::string_view(Arc::new(StringViewArray::from(vec!["hello view"])) as ArrayRef, DataType::Utf8View)]
|
||||
#[case::binary_view(Arc::new(BinaryViewArray::from(vec![b"\xCA\xFE".as_ref()])) as ArrayRef, DataType::BinaryView)]
|
||||
fn test_encode_decode_without_type_prefix(#[case] array: ArrayRef, #[case] dt: DataType) {
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
let opts = EncodeOptions {
|
||||
include_data_type: false,
|
||||
};
|
||||
let encoded = scalar.encode_with_options(&opts).unwrap();
|
||||
let decode_opts = DecodeOptions {
|
||||
data_type: Some(&dt),
|
||||
};
|
||||
let decoded = ArrowScalar::decode_with_options(&encoded, &decode_opts).unwrap();
|
||||
assert_eq!(scalar, decoded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_null_encode_decode_roundtrip() {
|
||||
let array: ArrayRef = Arc::new(Int32Array::from(vec![None]));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
assert!(scalar.is_null());
|
||||
let encoded = scalar.encode().unwrap();
|
||||
let decoded = ArrowScalar::decode(&encoded).unwrap();
|
||||
assert!(decoded.is_null());
|
||||
assert_eq!(decoded.data_type(), &DataType::Int32);
|
||||
assert_eq!(scalar, decoded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_null_encode_decode_without_type_prefix() {
|
||||
let array: ArrayRef = Arc::new(StringArray::from(vec![Option::<&str>::None]));
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
let opts = EncodeOptions {
|
||||
include_data_type: false,
|
||||
};
|
||||
let encoded = scalar.encode_with_options(&opts).unwrap();
|
||||
let decode_opts = DecodeOptions {
|
||||
data_type: Some(&DataType::Utf8),
|
||||
};
|
||||
let decoded = ArrowScalar::decode_with_options(&encoded, &decode_opts).unwrap();
|
||||
assert!(decoded.is_null());
|
||||
assert_eq!(decoded.data_type(), &DataType::Utf8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_trailing_bytes() {
|
||||
let scalar = ArrowScalar::from(42i32);
|
||||
let mut encoded = scalar.encode().unwrap();
|
||||
encoded.push(0xFF);
|
||||
assert!(ArrowScalar::decode(&encoded).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encoded_bytes_contain_format_prefix() {
|
||||
let scalar = ArrowScalar::from(42i32);
|
||||
let encoded = scalar.encode().unwrap();
|
||||
// First byte is varint length of format string "i" (length 1)
|
||||
assert_eq!(encoded[0], 1);
|
||||
// Second byte is the format string itself
|
||||
assert_eq!(encoded[1], b'i');
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::null(DataType::Null, "n")]
|
||||
#[case::boolean(DataType::Boolean, "b")]
|
||||
#[case::int8(DataType::Int8, "c")]
|
||||
#[case::uint8(DataType::UInt8, "C")]
|
||||
#[case::int16(DataType::Int16, "s")]
|
||||
#[case::uint16(DataType::UInt16, "S")]
|
||||
#[case::int32(DataType::Int32, "i")]
|
||||
#[case::uint32(DataType::UInt32, "I")]
|
||||
#[case::int64(DataType::Int64, "l")]
|
||||
#[case::uint64(DataType::UInt64, "L")]
|
||||
#[case::float16(DataType::Float16, "e")]
|
||||
#[case::float32(DataType::Float32, "f")]
|
||||
#[case::float64(DataType::Float64, "g")]
|
||||
#[case::binary(DataType::Binary, "z")]
|
||||
#[case::large_binary(DataType::LargeBinary, "Z")]
|
||||
#[case::utf8(DataType::Utf8, "u")]
|
||||
#[case::large_utf8(DataType::LargeUtf8, "U")]
|
||||
#[case::binary_view(DataType::BinaryView, "vz")]
|
||||
#[case::utf8_view(DataType::Utf8View, "vu")]
|
||||
#[case::date32(DataType::Date32, "tdD")]
|
||||
#[case::date64(DataType::Date64, "tdm")]
|
||||
#[case::fixed_size_binary(DataType::FixedSizeBinary(16), "w:16")]
|
||||
#[case::decimal128(DataType::Decimal128(10, 2), "d:10,2")]
|
||||
#[case::decimal256(DataType::Decimal256(38, 10), "d:38,10,256")]
|
||||
#[case::timestamp_us_utc(
|
||||
DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))),
|
||||
"tsu:UTC"
|
||||
)]
|
||||
#[case::timestamp_ns_none(DataType::Timestamp(TimeUnit::Nanosecond, None), "tsn:")]
|
||||
#[case::duration_s(DataType::Duration(TimeUnit::Second), "tDs")]
|
||||
#[case::interval_ym(DataType::Interval(IntervalUnit::YearMonth), "tiM")]
|
||||
fn test_format_string_roundtrip(#[case] dt: DataType, #[case] expected_fmt: &str) {
|
||||
let fmt = data_type_to_format_string(&dt).unwrap();
|
||||
assert_eq!(fmt.as_ref(), expected_fmt);
|
||||
let roundtripped = format_string_to_data_type(&fmt).unwrap();
|
||||
assert_eq!(roundtripped, dt);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_timestamp_with_tz_roundtrip() {
|
||||
let array: ArrayRef = Arc::new(
|
||||
TimestampMicrosecondArray::from(vec![1_000_000]).with_timezone("America/New_York"),
|
||||
);
|
||||
let scalar = ArrowScalar::try_from_array(array).unwrap();
|
||||
let encoded = scalar.encode().unwrap();
|
||||
let decoded = ArrowScalar::decode(&encoded).unwrap();
|
||||
assert_eq!(scalar, decoded);
|
||||
assert_eq!(scalar.data_type(), decoded.data_type());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
[package]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Statistics accumulator for Arrow arrays (min, max, null_count, nan_count)"
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
rust-version.workspace = true
|
||||
readme = "README.md"
|
||||
|
||||
[dependencies]
|
||||
arrow-array = { workspace = true }
|
||||
arrow-schema = { workspace = true }
|
||||
lance-arrow-scalar = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
arrow-select = { workspace = true }
|
||||
proptest = { workspace = true }
|
||||
rstest = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,62 @@
|
||||
# lance-arrow-stats
|
||||
|
||||
Statistics accumulator for [Apache Arrow](https://arrow.apache.org/) arrays.
|
||||
|
||||
Computes min, max, null count, NaN count, and buffer memory usage over one or
|
||||
more batches of Arrow data. Designed for use in Lance's columnar storage layer
|
||||
where page-level statistics drive predicate pushdown and query planning.
|
||||
|
||||
## Usage
|
||||
|
||||
```rust
|
||||
use arrow_array::{Int32Array, ArrayRef};
|
||||
use lance_arrow_stats::StatisticsAccumulator;
|
||||
use arrow_schema::DataType;
|
||||
use std::sync::Arc;
|
||||
|
||||
let mut acc = StatisticsAccumulator::new(&DataType::Int32);
|
||||
|
||||
let batch: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), None, Some(1), Some(4)]));
|
||||
acc.update(&batch).unwrap();
|
||||
|
||||
let stats = acc.finish();
|
||||
assert_eq!(stats.null_count, 1);
|
||||
```
|
||||
|
||||
## Tracked Statistics
|
||||
|
||||
| Statistic | Description |
|
||||
| --------------- | -------------------------------------------------------- |
|
||||
| `min` | Minimum non-null, non-NaN value (`ArrowScalar`) |
|
||||
| `max` | Maximum non-null, non-NaN value (`ArrowScalar`) |
|
||||
| `null_count` | Total number of null values |
|
||||
| `nan_count` | Total NaN values (float and float-list types only) |
|
||||
| `item_nulls` | Null items inside list entries (list types only) |
|
||||
| `buffer_memory` | Total Arrow buffer memory in bytes |
|
||||
|
||||
## Supported Types
|
||||
|
||||
- **Numeric** — Int8–Int64, UInt8–UInt64, Float16/32/64
|
||||
- **Temporal** — Date32/64, Time32/64, Timestamp, Duration
|
||||
- **Boolean**
|
||||
- **String** — Utf8, LargeUtf8
|
||||
- **Binary** — Binary, LargeBinary
|
||||
- **List** — List, LargeList, FixedSizeList (computes stats over items)
|
||||
|
||||
Dictionary, run-end encoded, and view types are accepted but min/max will be
|
||||
`None`.
|
||||
|
||||
## Merging
|
||||
|
||||
Accumulators of the same data type can be merged, which is useful for combining
|
||||
statistics computed in parallel across different pages or files:
|
||||
|
||||
```rust
|
||||
use lance_arrow_stats::StatisticsAccumulator;
|
||||
use arrow_schema::DataType;
|
||||
|
||||
let mut a = StatisticsAccumulator::new(&DataType::Float32);
|
||||
let mut b = StatisticsAccumulator::new(&DataType::Float32);
|
||||
// ... update each with different batches ...
|
||||
a.merge(&b).unwrap();
|
||||
```
|
||||
@@ -0,0 +1,8 @@
|
||||
# Seeds for failure cases proptest has generated in the past. It is
|
||||
# automatically read and these particular cases re-run before any
|
||||
# novel cases are generated.
|
||||
#
|
||||
# It is recommended to check this file in to source control so that
|
||||
# everyone who runs the test benefits from these saved cases.
|
||||
cc 81b0445f36fa8f491c1fb3162f51b61c8be140d5b2a1e792c42b4bdb7f1b6a62 # shrinks to values = [0.0, -0.0]
|
||||
cc 8651fce939497f33c6dafd842937d95965af97833bfbbd10df30d5ea00dbd07d # shrinks to values = [Some(0.0), Some(-0.0)]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,35 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use arrow_array::cast::AsArray;
|
||||
use arrow_array::types::{Float16Type, Float32Type, Float64Type};
|
||||
use arrow_array::{Array, ArrayRef};
|
||||
|
||||
macro_rules! count_nans_typed {
|
||||
($array:expr, $arrow_type:ty) => {{
|
||||
let typed = $array.as_primitive::<$arrow_type>();
|
||||
let mut count = 0u64;
|
||||
for i in 0..typed.len() {
|
||||
if !typed.is_null(i) && typed.value(i).is_nan() {
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
count
|
||||
}};
|
||||
}
|
||||
|
||||
/// Count the number of non-null NaN values in an array.
|
||||
///
|
||||
/// Returns 0 for non-float types.
|
||||
pub fn count_nans(array: &ArrayRef) -> u64 {
|
||||
use arrow_schema::DataType::*;
|
||||
match array.data_type() {
|
||||
Float16 => count_nans_typed!(array, Float16Type),
|
||||
Float32 => count_nans_typed!(array, Float32Type),
|
||||
Float64 => count_nans_typed!(array, Float64Type),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
[package]
|
||||
name = "lance-bitpacking"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
readme.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Vendored copy of https://github.com/spiraldb/fastlanes for use in Lance"
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
rust-version.workspace = true
|
||||
|
||||
[dependencies]
|
||||
arrayref = "0.3"
|
||||
crunchy = "0.2"
|
||||
paste = "1"
|
||||
seq-macro = "0.3"
|
||||
|
||||
[dev-dependencies]
|
||||
bitpacking = { workspace = true, features = ["bitpacker4x"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,823 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use super::BitPacker;
|
||||
|
||||
use crate::bitpacker_internal::{Available, UnsafeBitPacker};
|
||||
|
||||
const BLOCK_LEN: usize = 32 * 4;
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
mod sse3 {
|
||||
|
||||
use super::BLOCK_LEN;
|
||||
use crate::bitpacker_internal::Available;
|
||||
|
||||
use std::arch::x86_64::__m128i as DataType;
|
||||
use std::arch::x86_64::_mm_and_si128 as op_and;
|
||||
use std::arch::x86_64::_mm_lddqu_si128 as load_unaligned;
|
||||
use std::arch::x86_64::_mm_or_si128 as op_or;
|
||||
use std::arch::x86_64::_mm_set1_epi32 as set1;
|
||||
use std::arch::x86_64::_mm_slli_epi32 as left_shift_32;
|
||||
use std::arch::x86_64::_mm_srli_epi32 as right_shift_32;
|
||||
use std::arch::x86_64::_mm_storeu_si128 as store_unaligned;
|
||||
use std::arch::x86_64::{
|
||||
_mm_add_epi32, _mm_cvtsi128_si32, _mm_shuffle_epi32, _mm_slli_si128, _mm_srli_si128,
|
||||
_mm_sub_epi32,
|
||||
};
|
||||
|
||||
#[allow(non_snake_case)]
|
||||
#[inline]
|
||||
unsafe fn or_collapse_to_u32(accumulator: DataType) -> u32 {
|
||||
let a__b__c__d_ = accumulator;
|
||||
let ______a__b_ = _mm_srli_si128(a__b__c__d_, 8);
|
||||
let a__b__ca_db = op_or(a__b__c__d_, ______a__b_);
|
||||
let ___a__b__ca = _mm_srli_si128(a__b__ca_db, 4);
|
||||
let _______cadb = op_or(a__b__ca_db, ___a__b__ca);
|
||||
_mm_cvtsi128_si32(_______cadb) as u32
|
||||
}
|
||||
|
||||
#[target_feature(enable = "sse3")]
|
||||
unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType {
|
||||
_mm_sub_epi32(
|
||||
curr,
|
||||
op_or(_mm_slli_si128(curr, 4), _mm_srli_si128(prev, 12)),
|
||||
)
|
||||
}
|
||||
|
||||
#[target_feature(enable = "sse3")]
|
||||
#[allow(non_snake_case)]
|
||||
#[inline]
|
||||
unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType {
|
||||
let offset = _mm_shuffle_epi32(prev, 0xff);
|
||||
let a__b__c__d_ = delta;
|
||||
let ______a__b_ = _mm_slli_si128(delta, 8);
|
||||
let a__b__ca_db = _mm_add_epi32(______a__b_, a__b__c__d_);
|
||||
let ___a__b__ca = _mm_slli_si128(a__b__ca_db, 4);
|
||||
let a_ab_abc_abcd: DataType = _mm_add_epi32(___a__b__ca, a__b__ca_db);
|
||||
_mm_add_epi32(offset, a_ab_abc_abcd)
|
||||
}
|
||||
|
||||
#[target_feature(enable = "sse3")]
|
||||
#[inline]
|
||||
unsafe fn add(left: DataType, right: DataType) -> DataType {
|
||||
_mm_add_epi32(left, right)
|
||||
}
|
||||
|
||||
unsafe fn sub(left: DataType, right: DataType) -> DataType {
|
||||
_mm_sub_epi32(left, right)
|
||||
}
|
||||
|
||||
declare_bitpacker!(target_feature(enable = "sse3"));
|
||||
|
||||
impl Available for UnsafeBitPackerImpl {
|
||||
fn available() -> bool {
|
||||
is_x86_feature_detected!("sse3")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
mod neon {
|
||||
|
||||
use super::BLOCK_LEN;
|
||||
use crate::bitpacker_internal::Available;
|
||||
use std::arch::aarch64::{
|
||||
uint32x4_t, vaddq_u32, vandq_u32, vdupq_n_u32, vextq_u32, vgetq_lane_u32, vld1q_u32,
|
||||
vorrq_u32, vshlq_n_u32, vshrq_n_u32, vst1q_u32, vsubq_u32,
|
||||
};
|
||||
|
||||
pub(crate) type DataType = uint32x4_t;
|
||||
|
||||
#[inline]
|
||||
/// Creates a vector with all elements set to `el`.
|
||||
unsafe fn set1(el: i32) -> DataType {
|
||||
vdupq_n_u32(el as u32)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn right_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
const {
|
||||
assert!(N >= 0);
|
||||
assert!(N <= 32);
|
||||
}
|
||||
|
||||
// We unroll here because vshrq_n_u32 only accepts constants from 1 to 32.
|
||||
match N {
|
||||
0 => el,
|
||||
1 => vshrq_n_u32::<1>(el),
|
||||
2 => vshrq_n_u32::<2>(el),
|
||||
3 => vshrq_n_u32::<3>(el),
|
||||
4 => vshrq_n_u32::<4>(el),
|
||||
5 => vshrq_n_u32::<5>(el),
|
||||
6 => vshrq_n_u32::<6>(el),
|
||||
7 => vshrq_n_u32::<7>(el),
|
||||
8 => vshrq_n_u32::<8>(el),
|
||||
9 => vshrq_n_u32::<9>(el),
|
||||
10 => vshrq_n_u32::<10>(el),
|
||||
11 => vshrq_n_u32::<11>(el),
|
||||
12 => vshrq_n_u32::<12>(el),
|
||||
13 => vshrq_n_u32::<13>(el),
|
||||
14 => vshrq_n_u32::<14>(el),
|
||||
15 => vshrq_n_u32::<15>(el),
|
||||
16 => vshrq_n_u32::<16>(el),
|
||||
17 => vshrq_n_u32::<17>(el),
|
||||
18 => vshrq_n_u32::<18>(el),
|
||||
19 => vshrq_n_u32::<19>(el),
|
||||
20 => vshrq_n_u32::<20>(el),
|
||||
21 => vshrq_n_u32::<21>(el),
|
||||
22 => vshrq_n_u32::<22>(el),
|
||||
23 => vshrq_n_u32::<23>(el),
|
||||
24 => vshrq_n_u32::<24>(el),
|
||||
25 => vshrq_n_u32::<25>(el),
|
||||
26 => vshrq_n_u32::<26>(el),
|
||||
27 => vshrq_n_u32::<27>(el),
|
||||
28 => vshrq_n_u32::<28>(el),
|
||||
29 => vshrq_n_u32::<29>(el),
|
||||
30 => vshrq_n_u32::<30>(el),
|
||||
31 => vshrq_n_u32::<31>(el),
|
||||
32 => vdupq_n_u32(0),
|
||||
_ => core::hint::unreachable_unchecked(),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn left_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
const {
|
||||
assert!(N >= 0);
|
||||
assert!(N <= 32);
|
||||
}
|
||||
|
||||
// We unroll here because vshlq_n_u32 only accepts constants from 0 to 31.
|
||||
match N {
|
||||
0 => el,
|
||||
1 => vshlq_n_u32::<1>(el),
|
||||
2 => vshlq_n_u32::<2>(el),
|
||||
3 => vshlq_n_u32::<3>(el),
|
||||
4 => vshlq_n_u32::<4>(el),
|
||||
5 => vshlq_n_u32::<5>(el),
|
||||
6 => vshlq_n_u32::<6>(el),
|
||||
7 => vshlq_n_u32::<7>(el),
|
||||
8 => vshlq_n_u32::<8>(el),
|
||||
9 => vshlq_n_u32::<9>(el),
|
||||
10 => vshlq_n_u32::<10>(el),
|
||||
11 => vshlq_n_u32::<11>(el),
|
||||
12 => vshlq_n_u32::<12>(el),
|
||||
13 => vshlq_n_u32::<13>(el),
|
||||
14 => vshlq_n_u32::<14>(el),
|
||||
15 => vshlq_n_u32::<15>(el),
|
||||
16 => vshlq_n_u32::<16>(el),
|
||||
17 => vshlq_n_u32::<17>(el),
|
||||
18 => vshlq_n_u32::<18>(el),
|
||||
19 => vshlq_n_u32::<19>(el),
|
||||
20 => vshlq_n_u32::<20>(el),
|
||||
21 => vshlq_n_u32::<21>(el),
|
||||
22 => vshlq_n_u32::<22>(el),
|
||||
23 => vshlq_n_u32::<23>(el),
|
||||
24 => vshlq_n_u32::<24>(el),
|
||||
25 => vshlq_n_u32::<25>(el),
|
||||
26 => vshlq_n_u32::<26>(el),
|
||||
27 => vshlq_n_u32::<27>(el),
|
||||
28 => vshlq_n_u32::<28>(el),
|
||||
29 => vshlq_n_u32::<29>(el),
|
||||
30 => vshlq_n_u32::<30>(el),
|
||||
31 => vshlq_n_u32::<31>(el),
|
||||
32 => vdupq_n_u32(0),
|
||||
_ => core::hint::unreachable_unchecked(),
|
||||
}
|
||||
}
|
||||
|
||||
use vorrq_u32 as op_or;
|
||||
|
||||
#[inline]
|
||||
unsafe fn op_and(left: DataType, right: DataType) -> DataType {
|
||||
vandq_u32(left, right)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn load_unaligned(addr: *const DataType) -> DataType {
|
||||
vld1q_u32(addr.cast::<u32>())
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn store_unaligned(addr: *mut DataType, data: DataType) {
|
||||
vst1q_u32(addr.cast::<u32>(), data);
|
||||
}
|
||||
|
||||
#[inline]
|
||||
/// Collapses the vector by performing a bitwise OR across all lanes
|
||||
unsafe fn or_collapse_to_u32(acc: DataType) -> u32 {
|
||||
vgetq_lane_u32(acc, 0)
|
||||
| vgetq_lane_u32(acc, 1)
|
||||
| vgetq_lane_u32(acc, 2)
|
||||
| vgetq_lane_u32(acc, 3)
|
||||
}
|
||||
|
||||
unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType {
|
||||
// Build a vector with [prev[3], curr[0], curr[1], curr[2]]
|
||||
let prev_shifted = vextq_u32(prev, curr, 3);
|
||||
vsubq_u32(curr, prev_shifted)
|
||||
}
|
||||
|
||||
#[allow(non_snake_case)]
|
||||
#[inline]
|
||||
unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType {
|
||||
let base = vdupq_n_u32(vgetq_lane_u32(prev, 3));
|
||||
let zero = vdupq_n_u32(0);
|
||||
let a__b__c__d_ = delta;
|
||||
let ______a__b_ = vextq_u32(zero, a__b__c__d_, 2);
|
||||
let a__b__ca_db = vaddq_u32(______a__b_, a__b__c__d_);
|
||||
let ___a__b__ca = vextq_u32(zero, a__b__ca_db, 3);
|
||||
let a_ab_abc_abcd = vaddq_u32(___a__b__ca, a__b__ca_db);
|
||||
vaddq_u32(base, a_ab_abc_abcd)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn add(left: DataType, right: DataType) -> DataType {
|
||||
vaddq_u32(left, right)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn sub(left: DataType, right: DataType) -> DataType {
|
||||
vsubq_u32(left, right)
|
||||
}
|
||||
|
||||
declare_bitpacker!(target_feature(enable = "neon"));
|
||||
|
||||
impl Available for UnsafeBitPackerImpl {
|
||||
fn available() -> bool {
|
||||
std::arch::is_aarch64_feature_detected!("neon")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mod scalar {
|
||||
|
||||
use super::BLOCK_LEN;
|
||||
use crate::bitpacker_internal::Available;
|
||||
use std::ptr;
|
||||
|
||||
pub(crate) type DataType = [u32; 4];
|
||||
|
||||
pub(crate) fn set1(el: i32) -> DataType {
|
||||
[el as u32; 4]
|
||||
}
|
||||
|
||||
pub(crate) fn right_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
[el[0] >> N, el[1] >> N, el[2] >> N, el[3] >> N]
|
||||
}
|
||||
|
||||
pub(crate) fn left_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
[el[0] << N, el[1] << N, el[2] << N, el[3] << N]
|
||||
}
|
||||
|
||||
pub(crate) fn op_or(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0] | right[0],
|
||||
left[1] | right[1],
|
||||
left[2] | right[2],
|
||||
left[3] | right[3],
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) fn op_and(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0] & right[0],
|
||||
left[1] & right[1],
|
||||
left[2] & right[2],
|
||||
left[3] & right[3],
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) unsafe fn load_unaligned(addr: *const DataType) -> DataType {
|
||||
ptr::read_unaligned(addr)
|
||||
}
|
||||
|
||||
pub(crate) unsafe fn store_unaligned(addr: *mut DataType, data: DataType) {
|
||||
ptr::write_unaligned(addr, data);
|
||||
}
|
||||
|
||||
pub(crate) fn or_collapse_to_u32(accumulator: DataType) -> u32 {
|
||||
(accumulator[0] | accumulator[1]) | (accumulator[2] | accumulator[3])
|
||||
}
|
||||
|
||||
fn compute_delta(curr: DataType, prev: DataType) -> DataType {
|
||||
[
|
||||
curr[0].wrapping_sub(prev[3]),
|
||||
curr[1].wrapping_sub(curr[0]),
|
||||
curr[2].wrapping_sub(curr[1]),
|
||||
curr[3].wrapping_sub(curr[2]),
|
||||
]
|
||||
}
|
||||
|
||||
fn integrate_delta(offset: DataType, delta: DataType) -> DataType {
|
||||
let el0 = offset[3].wrapping_add(delta[0]);
|
||||
let el1 = el0.wrapping_add(delta[1]);
|
||||
let el2 = el1.wrapping_add(delta[2]);
|
||||
let el3 = el2.wrapping_add(delta[3]);
|
||||
[el0, el1, el2, el3]
|
||||
}
|
||||
|
||||
pub(crate) fn add(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0].wrapping_add(right[0]),
|
||||
left[1].wrapping_add(right[1]),
|
||||
left[2].wrapping_add(right[2]),
|
||||
left[3].wrapping_add(right[3]),
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) fn sub(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0].wrapping_sub(right[0]),
|
||||
left[1].wrapping_sub(right[1]),
|
||||
left[2].wrapping_sub(right[2]),
|
||||
left[3].wrapping_sub(right[3]),
|
||||
]
|
||||
}
|
||||
|
||||
// The `allow(unused)` is here to put an attribute that has no effect.
|
||||
//
|
||||
// For other bitpacker, we enable specific CPU instruction set, but for the
|
||||
// scalar bitpacker none is required.
|
||||
declare_bitpacker!(allow(unused));
|
||||
|
||||
impl Available for UnsafeBitPackerImpl {
|
||||
fn available() -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum InstructionSet {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
SSE3,
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
NEON,
|
||||
Scalar,
|
||||
}
|
||||
|
||||
/// `BitPacker4x` packs integers in groups of 4. This gives an opportunity
|
||||
/// to leverage `SSE3` instructions to encode and decode the stream.
|
||||
///
|
||||
/// One block must contain `128 integers`.
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct BitPacker4x(InstructionSet);
|
||||
|
||||
impl BitPacker4x {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
pub(crate) fn new_sse() -> Option<Self> {
|
||||
sse3::UnsafeBitPackerImpl::available().then_some(BitPacker4x(InstructionSet::SSE3))
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "x86_64"))]
|
||||
pub(crate) fn new_sse() -> Option<Self> {
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
pub(crate) fn new_neon() -> Option<Self> {
|
||||
neon::UnsafeBitPackerImpl::available().then_some(BitPacker4x(InstructionSet::NEON))
|
||||
}
|
||||
|
||||
#[cfg(not(all(target_arch = "aarch64", target_endian = "little")))]
|
||||
pub(crate) fn new_neon() -> Option<Self> {
|
||||
None
|
||||
}
|
||||
|
||||
pub(crate) fn new_scalar() -> Self {
|
||||
BitPacker4x(InstructionSet::Scalar)
|
||||
}
|
||||
}
|
||||
|
||||
impl BitPacker for BitPacker4x {
|
||||
const BLOCK_LEN: usize = BLOCK_LEN;
|
||||
|
||||
fn new() -> Self {
|
||||
Self::new_sse()
|
||||
.or_else(Self::new_neon)
|
||||
.unwrap_or_else(Self::new_scalar)
|
||||
}
|
||||
|
||||
fn compress(&self, decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => {
|
||||
sse3::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn compress_sorted(
|
||||
&self,
|
||||
initial: u32,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::compress_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn compress_strictly_sorted(
|
||||
&self,
|
||||
initial: Option<u32>,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::compress_strictly_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_strictly_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_strictly_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decompress(&self, compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => {
|
||||
sse3::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decompress_strictly_sorted(
|
||||
&self,
|
||||
initial: Option<u32>,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::decompress_strictly_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_strictly_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_strictly_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decompress_sorted(
|
||||
&self,
|
||||
initial: u32,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::decompress_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn num_bits(&self, decompressed: &[u32]) -> u8 {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => sse3::UnsafeBitPackerImpl::num_bits(decompressed),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::num_bits(decompressed),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::num_bits(decompressed),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn num_bits_sorted(&self, initial: u32, decompressed: &[u32]) -> u8 {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => {
|
||||
sse3::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn num_bits_strictly_sorted(&self, initial: Option<u32>, decompressed: &[u32]) -> u8 {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::SSE3 => {
|
||||
sse3::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_arch = "x86_64",
|
||||
all(target_arch = "aarch64", target_endian = "little")
|
||||
))]
|
||||
#[cfg(any())]
|
||||
mod tests {
|
||||
use super::BLOCK_LEN;
|
||||
use super::scalar;
|
||||
use crate::bitpacker_internal::Available;
|
||||
use crate::tests::test_util_compatible;
|
||||
use crate::{BitPacker, BitPacker4x};
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
#[test]
|
||||
fn test_compatible_sse3() {
|
||||
use super::sse3;
|
||||
if sse3::UnsafeBitPackerImpl::available() {
|
||||
test_util_compatible::<scalar::UnsafeBitPackerImpl, sse3::UnsafeBitPackerImpl>(
|
||||
BLOCK_LEN,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
#[test]
|
||||
fn test_compatible_neon() {
|
||||
use super::neon;
|
||||
if neon::UnsafeBitPackerImpl::available() {
|
||||
test_util_compatible::<scalar::UnsafeBitPackerImpl, neon::UnsafeBitPackerImpl>(
|
||||
BLOCK_LEN,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_bit_width_32() {
|
||||
let values = vec![i32::max_value() as u32 + 1; BitPacker4x::BLOCK_LEN];
|
||||
let bit_packer = BitPacker4x::new();
|
||||
let bit_width = bit_packer.num_bits_sorted(0, &values);
|
||||
assert_eq!(bit_width, 32);
|
||||
|
||||
let mut block = vec![0u8; BitPacker4x::compressed_block_size(bit_width)];
|
||||
bit_packer.compress_sorted(0, &values, &mut block, bit_width);
|
||||
|
||||
let mut decoded_values = vec![0x10101010; BitPacker4x::BLOCK_LEN];
|
||||
bit_packer.decompress_sorted(0, &block, &mut decoded_values, bit_width);
|
||||
|
||||
assert_eq!(values, decoded_values);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bit_width_32() {
|
||||
let mut values = vec![i32::max_value() as u32 + 1; BitPacker4x::BLOCK_LEN];
|
||||
values[0] = 0;
|
||||
let bit_packer = BitPacker4x::new();
|
||||
let bit_width = bit_packer.num_bits(&values);
|
||||
assert_eq!(bit_width, 32);
|
||||
|
||||
let mut block = vec![0u8; BitPacker4x::compressed_block_size(bit_width)];
|
||||
bit_packer.compress(&values, &mut block, bit_width);
|
||||
|
||||
let mut decoded_values = vec![0x10101010; BitPacker4x::BLOCK_LEN];
|
||||
bit_packer.decompress(&block, &mut decoded_values, bit_width);
|
||||
|
||||
assert_eq!(values, decoded_values);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{BLOCK_LEN, BitPacker4x};
|
||||
use crate::bitpacker_internal::BitPacker;
|
||||
use bitpacking::{BitPacker as ExternalBitPacker, BitPacker4x as ExternalBitPacker4x};
|
||||
|
||||
fn mask_for_width(width: u8) -> u32 {
|
||||
match width {
|
||||
0 => 0,
|
||||
32 => u32::MAX,
|
||||
_ => (1u32 << width) - 1,
|
||||
}
|
||||
}
|
||||
|
||||
fn raw_values(width: u8) -> Vec<u32> {
|
||||
let mask = mask_for_width(width);
|
||||
(0..BLOCK_LEN)
|
||||
.map(|idx| ((idx * 17 + 3) as u32) & mask)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sorted_values(width: u8) -> (u32, Vec<u32>) {
|
||||
if width == 0 {
|
||||
return (11, vec![11; BLOCK_LEN]);
|
||||
}
|
||||
if width == 32 {
|
||||
return (0, vec![u32::MAX; BLOCK_LEN]);
|
||||
}
|
||||
|
||||
let mask = mask_for_width(width).min(127);
|
||||
let initial = 11u32;
|
||||
let mut current = initial;
|
||||
let values = (0..BLOCK_LEN)
|
||||
.map(|idx| {
|
||||
current += (idx as u32 * 7 + 1) & mask;
|
||||
current
|
||||
})
|
||||
.collect();
|
||||
(initial, values)
|
||||
}
|
||||
|
||||
fn strictly_sorted_values(width: u8) -> (Option<u32>, Vec<u32>) {
|
||||
let mask = mask_for_width(width).min(127);
|
||||
let mut current = 0u32;
|
||||
let values = (0..BLOCK_LEN)
|
||||
.map(|idx| {
|
||||
if idx == 0 {
|
||||
current = 0;
|
||||
} else {
|
||||
current += 1 + ((idx as u32 * 5) & mask);
|
||||
}
|
||||
current
|
||||
})
|
||||
.collect();
|
||||
(None, values)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scalar_backend_matches_external_bitpacker4x() {
|
||||
let scalar = BitPacker4x::new_scalar();
|
||||
let external = ExternalBitPacker4x::new();
|
||||
|
||||
for width in 0..=32 {
|
||||
let values = raw_values(width);
|
||||
assert_eq!(scalar.num_bits(&values), external.num_bits(&values));
|
||||
|
||||
let mut actual = vec![0u8; BitPacker4x::compressed_block_size(width)];
|
||||
let actual_len = scalar.compress(&values, &mut actual, width);
|
||||
|
||||
let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(width)];
|
||||
let expected_len = external.compress(&values, &mut expected, width);
|
||||
|
||||
assert_eq!(actual_len, expected_len);
|
||||
assert_eq!(actual, expected, "raw width {width}");
|
||||
|
||||
let mut decoded = vec![0u32; BLOCK_LEN];
|
||||
assert_eq!(scalar.decompress(&actual, &mut decoded, width), actual_len);
|
||||
assert_eq!(decoded, values);
|
||||
|
||||
let (initial, values) = sorted_values(width);
|
||||
assert_eq!(
|
||||
scalar.num_bits_sorted(initial, &values),
|
||||
external.num_bits_sorted(initial, &values)
|
||||
);
|
||||
|
||||
let mut actual = vec![0u8; BitPacker4x::compressed_block_size(width)];
|
||||
let actual_len = scalar.compress_sorted(initial, &values, &mut actual, width);
|
||||
|
||||
let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(width)];
|
||||
let expected_len = external.compress_sorted(initial, &values, &mut expected, width);
|
||||
|
||||
assert_eq!(actual_len, expected_len);
|
||||
assert_eq!(actual, expected, "sorted width {width}");
|
||||
|
||||
let mut decoded = vec![0u32; BLOCK_LEN];
|
||||
assert_eq!(
|
||||
scalar.decompress_sorted(initial, &actual, &mut decoded, width),
|
||||
actual_len
|
||||
);
|
||||
assert_eq!(decoded, values);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scalar_backend_matches_external_strictly_sorted_bitpacker4x() {
|
||||
let scalar = BitPacker4x::new_scalar();
|
||||
let external = ExternalBitPacker4x::new();
|
||||
|
||||
for width in 0..=16 {
|
||||
let (initial, values) = strictly_sorted_values(width);
|
||||
let num_bits = external.num_bits_strictly_sorted(initial, &values);
|
||||
assert_eq!(scalar.num_bits_strictly_sorted(initial, &values), num_bits);
|
||||
|
||||
let mut actual = vec![0u8; BitPacker4x::compressed_block_size(num_bits)];
|
||||
let actual_len =
|
||||
scalar.compress_strictly_sorted(initial, &values, &mut actual, num_bits);
|
||||
|
||||
let mut expected = vec![0u8; ExternalBitPacker4x::compressed_block_size(num_bits)];
|
||||
let expected_len =
|
||||
external.compress_strictly_sorted(initial, &values, &mut expected, num_bits);
|
||||
|
||||
assert_eq!(actual_len, expected_len);
|
||||
assert_eq!(actual, expected, "strict width {width}");
|
||||
|
||||
let mut decoded = vec![0u32; BLOCK_LEN];
|
||||
assert_eq!(
|
||||
scalar.decompress_strictly_sorted(initial, &actual, &mut decoded, num_bits),
|
||||
actual_len
|
||||
);
|
||||
assert_eq!(decoded, values);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,875 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use super::BitPacker;
|
||||
|
||||
use crate::bitpacker_internal::{Available, UnsafeBitPacker};
|
||||
|
||||
const BLOCK_LEN: usize = 32 * 8;
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
mod avx2 {
|
||||
use super::BLOCK_LEN;
|
||||
use crate::bitpacker_internal::Available;
|
||||
|
||||
use std::arch::x86_64::__m256i as DataType;
|
||||
use std::arch::x86_64::_mm256_and_si256 as op_and;
|
||||
use std::arch::x86_64::_mm256_lddqu_si256 as load_unaligned;
|
||||
use std::arch::x86_64::_mm256_or_si256 as op_or;
|
||||
use std::arch::x86_64::_mm256_set1_epi32 as set1;
|
||||
use std::arch::x86_64::_mm256_slli_epi32 as left_shift_32;
|
||||
use std::arch::x86_64::_mm256_srli_epi32 as right_shift_32;
|
||||
use std::arch::x86_64::_mm256_storeu_si256 as store_unaligned;
|
||||
|
||||
use std::arch::x86_64::{
|
||||
_mm256_add_epi32, _mm256_extract_epi32, _mm256_permute2f128_si256, _mm256_shuffle_epi32,
|
||||
_mm256_slli_si256, _mm256_srli_si256, _mm256_sub_epi32,
|
||||
};
|
||||
|
||||
#[allow(non_snake_case)]
|
||||
unsafe fn or_collapse_to_u32(accumulator: DataType) -> u32 {
|
||||
let a__b__c__d__e__f__g__h_ = accumulator;
|
||||
let ______a__b________e__f = _mm256_srli_si256(a__b__c__d__e__f__g__h_, 8);
|
||||
let a__b__ca_db_e__f__ge_hf = op_or(a__b__c__d__e__f__g__h_, ______a__b________e__f);
|
||||
let ___a__b__ca____e__f__ge = _mm256_srli_si256(a__b__ca_db_e__f__ge_hf, 4);
|
||||
let _________cadb______gehf = op_or(a__b__ca_db_e__f__ge_hf, ___a__b__ca____e__f__ge);
|
||||
let cadb = _mm256_extract_epi32(_________cadb______gehf, 0);
|
||||
let gehf = _mm256_extract_epi32(_________cadb______gehf, 4);
|
||||
(cadb | gehf) as u32
|
||||
}
|
||||
|
||||
unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType {
|
||||
let left_shift = _mm256_slli_si256(curr, 4);
|
||||
let curr_shift = _mm256_srli_si256(curr, 12);
|
||||
let curr_right_only = _mm256_permute2f128_si256(curr_shift, curr_shift, 8);
|
||||
let prev_shift = _mm256_srli_si256(prev, 12);
|
||||
let sub_left = _mm256_permute2f128_si256(prev_shift, prev_shift, 3 | (8 << 4));
|
||||
let diff = op_or(left_shift, op_or(curr_right_only, sub_left));
|
||||
_mm256_sub_epi32(curr, diff)
|
||||
}
|
||||
|
||||
#[allow(non_snake_case)]
|
||||
unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType {
|
||||
let offset_repeat = _mm256_shuffle_epi32(prev, 0xff);
|
||||
let offset = _mm256_permute2f128_si256(offset_repeat, offset_repeat, 3 | (8 << 4));
|
||||
let a__b__c__d__e__f__g__h__ = delta;
|
||||
let ______a__b________e__f__ = _mm256_slli_si256(delta, 8);
|
||||
let a__b__ca_db_e__f__ge_fh_ =
|
||||
_mm256_add_epi32(a__b__c__d__e__f__g__h__, ______a__b________e__f__);
|
||||
let ___a__b__ca____e__f__ge_ = _mm256_slli_si256(a__b__ca_db_e__f__ge_fh_, 4);
|
||||
let halved_prefix_sum =
|
||||
_mm256_add_epi32(___a__b__ca____e__f__ge_, a__b__ca_db_e__f__ge_fh_);
|
||||
let offsetted_halved_prefix_sum = _mm256_add_epi32(halved_prefix_sum, offset);
|
||||
let select_last_low = _mm256_shuffle_epi32(offsetted_halved_prefix_sum, 0xff);
|
||||
let high_offset = _mm256_permute2f128_si256(select_last_low, select_last_low, 8);
|
||||
_mm256_add_epi32(high_offset, offsetted_halved_prefix_sum)
|
||||
}
|
||||
|
||||
unsafe fn add(left: DataType, right: DataType) -> DataType {
|
||||
_mm256_add_epi32(left, right)
|
||||
}
|
||||
|
||||
unsafe fn sub(left: DataType, right: DataType) -> DataType {
|
||||
_mm256_sub_epi32(left, right)
|
||||
}
|
||||
|
||||
declare_bitpacker!(target_feature(enable = "avx2"));
|
||||
|
||||
impl Available for UnsafeBitPackerImpl {
|
||||
fn available() -> bool {
|
||||
is_x86_feature_detected!("avx2")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
mod neon {
|
||||
use super::BLOCK_LEN;
|
||||
use crate::bitpacker_internal::Available;
|
||||
use std::arch::aarch64::{
|
||||
uint32x4_t, vaddq_u32, vandq_u32, vdupq_n_u32, vextq_u32, vgetq_lane_u32, vld1q_u32,
|
||||
vorrq_u32, vshlq_n_u32, vshrq_n_u32, vst1q_u32, vsubq_u32,
|
||||
};
|
||||
|
||||
pub(crate) type DataType = [uint32x4_t; 2];
|
||||
|
||||
#[inline]
|
||||
unsafe fn set1(el: i32) -> DataType {
|
||||
let lanes = vdupq_n_u32(el as u32);
|
||||
[lanes, lanes]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn right_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
const {
|
||||
assert!(N >= 0);
|
||||
assert!(N <= 32);
|
||||
}
|
||||
|
||||
match N {
|
||||
0 => el,
|
||||
1 => [vshrq_n_u32::<1>(el[0]), vshrq_n_u32::<1>(el[1])],
|
||||
2 => [vshrq_n_u32::<2>(el[0]), vshrq_n_u32::<2>(el[1])],
|
||||
3 => [vshrq_n_u32::<3>(el[0]), vshrq_n_u32::<3>(el[1])],
|
||||
4 => [vshrq_n_u32::<4>(el[0]), vshrq_n_u32::<4>(el[1])],
|
||||
5 => [vshrq_n_u32::<5>(el[0]), vshrq_n_u32::<5>(el[1])],
|
||||
6 => [vshrq_n_u32::<6>(el[0]), vshrq_n_u32::<6>(el[1])],
|
||||
7 => [vshrq_n_u32::<7>(el[0]), vshrq_n_u32::<7>(el[1])],
|
||||
8 => [vshrq_n_u32::<8>(el[0]), vshrq_n_u32::<8>(el[1])],
|
||||
9 => [vshrq_n_u32::<9>(el[0]), vshrq_n_u32::<9>(el[1])],
|
||||
10 => [vshrq_n_u32::<10>(el[0]), vshrq_n_u32::<10>(el[1])],
|
||||
11 => [vshrq_n_u32::<11>(el[0]), vshrq_n_u32::<11>(el[1])],
|
||||
12 => [vshrq_n_u32::<12>(el[0]), vshrq_n_u32::<12>(el[1])],
|
||||
13 => [vshrq_n_u32::<13>(el[0]), vshrq_n_u32::<13>(el[1])],
|
||||
14 => [vshrq_n_u32::<14>(el[0]), vshrq_n_u32::<14>(el[1])],
|
||||
15 => [vshrq_n_u32::<15>(el[0]), vshrq_n_u32::<15>(el[1])],
|
||||
16 => [vshrq_n_u32::<16>(el[0]), vshrq_n_u32::<16>(el[1])],
|
||||
17 => [vshrq_n_u32::<17>(el[0]), vshrq_n_u32::<17>(el[1])],
|
||||
18 => [vshrq_n_u32::<18>(el[0]), vshrq_n_u32::<18>(el[1])],
|
||||
19 => [vshrq_n_u32::<19>(el[0]), vshrq_n_u32::<19>(el[1])],
|
||||
20 => [vshrq_n_u32::<20>(el[0]), vshrq_n_u32::<20>(el[1])],
|
||||
21 => [vshrq_n_u32::<21>(el[0]), vshrq_n_u32::<21>(el[1])],
|
||||
22 => [vshrq_n_u32::<22>(el[0]), vshrq_n_u32::<22>(el[1])],
|
||||
23 => [vshrq_n_u32::<23>(el[0]), vshrq_n_u32::<23>(el[1])],
|
||||
24 => [vshrq_n_u32::<24>(el[0]), vshrq_n_u32::<24>(el[1])],
|
||||
25 => [vshrq_n_u32::<25>(el[0]), vshrq_n_u32::<25>(el[1])],
|
||||
26 => [vshrq_n_u32::<26>(el[0]), vshrq_n_u32::<26>(el[1])],
|
||||
27 => [vshrq_n_u32::<27>(el[0]), vshrq_n_u32::<27>(el[1])],
|
||||
28 => [vshrq_n_u32::<28>(el[0]), vshrq_n_u32::<28>(el[1])],
|
||||
29 => [vshrq_n_u32::<29>(el[0]), vshrq_n_u32::<29>(el[1])],
|
||||
30 => [vshrq_n_u32::<30>(el[0]), vshrq_n_u32::<30>(el[1])],
|
||||
31 => [vshrq_n_u32::<31>(el[0]), vshrq_n_u32::<31>(el[1])],
|
||||
32 => set1(0),
|
||||
_ => core::hint::unreachable_unchecked(),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn left_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
const {
|
||||
assert!(N >= 0);
|
||||
assert!(N <= 32);
|
||||
}
|
||||
|
||||
match N {
|
||||
0 => el,
|
||||
1 => [vshlq_n_u32::<1>(el[0]), vshlq_n_u32::<1>(el[1])],
|
||||
2 => [vshlq_n_u32::<2>(el[0]), vshlq_n_u32::<2>(el[1])],
|
||||
3 => [vshlq_n_u32::<3>(el[0]), vshlq_n_u32::<3>(el[1])],
|
||||
4 => [vshlq_n_u32::<4>(el[0]), vshlq_n_u32::<4>(el[1])],
|
||||
5 => [vshlq_n_u32::<5>(el[0]), vshlq_n_u32::<5>(el[1])],
|
||||
6 => [vshlq_n_u32::<6>(el[0]), vshlq_n_u32::<6>(el[1])],
|
||||
7 => [vshlq_n_u32::<7>(el[0]), vshlq_n_u32::<7>(el[1])],
|
||||
8 => [vshlq_n_u32::<8>(el[0]), vshlq_n_u32::<8>(el[1])],
|
||||
9 => [vshlq_n_u32::<9>(el[0]), vshlq_n_u32::<9>(el[1])],
|
||||
10 => [vshlq_n_u32::<10>(el[0]), vshlq_n_u32::<10>(el[1])],
|
||||
11 => [vshlq_n_u32::<11>(el[0]), vshlq_n_u32::<11>(el[1])],
|
||||
12 => [vshlq_n_u32::<12>(el[0]), vshlq_n_u32::<12>(el[1])],
|
||||
13 => [vshlq_n_u32::<13>(el[0]), vshlq_n_u32::<13>(el[1])],
|
||||
14 => [vshlq_n_u32::<14>(el[0]), vshlq_n_u32::<14>(el[1])],
|
||||
15 => [vshlq_n_u32::<15>(el[0]), vshlq_n_u32::<15>(el[1])],
|
||||
16 => [vshlq_n_u32::<16>(el[0]), vshlq_n_u32::<16>(el[1])],
|
||||
17 => [vshlq_n_u32::<17>(el[0]), vshlq_n_u32::<17>(el[1])],
|
||||
18 => [vshlq_n_u32::<18>(el[0]), vshlq_n_u32::<18>(el[1])],
|
||||
19 => [vshlq_n_u32::<19>(el[0]), vshlq_n_u32::<19>(el[1])],
|
||||
20 => [vshlq_n_u32::<20>(el[0]), vshlq_n_u32::<20>(el[1])],
|
||||
21 => [vshlq_n_u32::<21>(el[0]), vshlq_n_u32::<21>(el[1])],
|
||||
22 => [vshlq_n_u32::<22>(el[0]), vshlq_n_u32::<22>(el[1])],
|
||||
23 => [vshlq_n_u32::<23>(el[0]), vshlq_n_u32::<23>(el[1])],
|
||||
24 => [vshlq_n_u32::<24>(el[0]), vshlq_n_u32::<24>(el[1])],
|
||||
25 => [vshlq_n_u32::<25>(el[0]), vshlq_n_u32::<25>(el[1])],
|
||||
26 => [vshlq_n_u32::<26>(el[0]), vshlq_n_u32::<26>(el[1])],
|
||||
27 => [vshlq_n_u32::<27>(el[0]), vshlq_n_u32::<27>(el[1])],
|
||||
28 => [vshlq_n_u32::<28>(el[0]), vshlq_n_u32::<28>(el[1])],
|
||||
29 => [vshlq_n_u32::<29>(el[0]), vshlq_n_u32::<29>(el[1])],
|
||||
30 => [vshlq_n_u32::<30>(el[0]), vshlq_n_u32::<30>(el[1])],
|
||||
31 => [vshlq_n_u32::<31>(el[0]), vshlq_n_u32::<31>(el[1])],
|
||||
32 => set1(0),
|
||||
_ => core::hint::unreachable_unchecked(),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn op_or(left: DataType, right: DataType) -> DataType {
|
||||
[vorrq_u32(left[0], right[0]), vorrq_u32(left[1], right[1])]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn op_and(left: DataType, right: DataType) -> DataType {
|
||||
[vandq_u32(left[0], right[0]), vandq_u32(left[1], right[1])]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn load_unaligned(addr: *const DataType) -> DataType {
|
||||
let ptr = addr.cast::<u32>();
|
||||
[vld1q_u32(ptr), vld1q_u32(ptr.add(4))]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn store_unaligned(addr: *mut DataType, data: DataType) {
|
||||
let ptr = addr.cast::<u32>();
|
||||
vst1q_u32(ptr, data[0]);
|
||||
vst1q_u32(ptr.add(4), data[1]);
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn or_collapse_to_u32(accumulator: DataType) -> u32 {
|
||||
vgetq_lane_u32(accumulator[0], 0)
|
||||
| vgetq_lane_u32(accumulator[0], 1)
|
||||
| vgetq_lane_u32(accumulator[0], 2)
|
||||
| vgetq_lane_u32(accumulator[0], 3)
|
||||
| vgetq_lane_u32(accumulator[1], 0)
|
||||
| vgetq_lane_u32(accumulator[1], 1)
|
||||
| vgetq_lane_u32(accumulator[1], 2)
|
||||
| vgetq_lane_u32(accumulator[1], 3)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn compute_delta(curr: DataType, prev: DataType) -> DataType {
|
||||
[
|
||||
vsubq_u32(curr[0], vextq_u32(prev[1], curr[0], 3)),
|
||||
vsubq_u32(curr[1], vextq_u32(curr[0], curr[1], 3)),
|
||||
]
|
||||
}
|
||||
|
||||
#[allow(non_snake_case)]
|
||||
#[inline]
|
||||
unsafe fn integrate_half(base: u32, delta: uint32x4_t) -> uint32x4_t {
|
||||
let base = vdupq_n_u32(base);
|
||||
let zero = vdupq_n_u32(0);
|
||||
let a__b__c__d_ = delta;
|
||||
let ______a__b_ = vextq_u32(zero, a__b__c__d_, 2);
|
||||
let a__b__ca_db = vaddq_u32(______a__b_, a__b__c__d_);
|
||||
let ___a__b__ca = vextq_u32(zero, a__b__ca_db, 3);
|
||||
let a_ab_abc_abcd = vaddq_u32(___a__b__ca, a__b__ca_db);
|
||||
vaddq_u32(base, a_ab_abc_abcd)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn integrate_delta(prev: DataType, delta: DataType) -> DataType {
|
||||
let low = integrate_half(vgetq_lane_u32(prev[1], 3), delta[0]);
|
||||
let high = integrate_half(vgetq_lane_u32(low, 3), delta[1]);
|
||||
[low, high]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn add(left: DataType, right: DataType) -> DataType {
|
||||
[vaddq_u32(left[0], right[0]), vaddq_u32(left[1], right[1])]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn sub(left: DataType, right: DataType) -> DataType {
|
||||
[vsubq_u32(left[0], right[0]), vsubq_u32(left[1], right[1])]
|
||||
}
|
||||
|
||||
declare_bitpacker!(target_feature(enable = "neon"));
|
||||
|
||||
impl Available for UnsafeBitPackerImpl {
|
||||
fn available() -> bool {
|
||||
std::arch::is_aarch64_feature_detected!("neon")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mod scalar {
|
||||
use super::BLOCK_LEN;
|
||||
use crate::bitpacker_internal::Available;
|
||||
use std::ptr;
|
||||
|
||||
pub(crate) type DataType = [u32; 8];
|
||||
|
||||
pub(crate) fn set1(el: i32) -> DataType {
|
||||
[el as u32; 8]
|
||||
}
|
||||
|
||||
pub(crate) fn right_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
[
|
||||
el[0] >> N,
|
||||
el[1] >> N,
|
||||
el[2] >> N,
|
||||
el[3] >> N,
|
||||
el[4] >> N,
|
||||
el[5] >> N,
|
||||
el[6] >> N,
|
||||
el[7] >> N,
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) fn left_shift_32<const N: i32>(el: DataType) -> DataType {
|
||||
[
|
||||
el[0] << N,
|
||||
el[1] << N,
|
||||
el[2] << N,
|
||||
el[3] << N,
|
||||
el[4] << N,
|
||||
el[5] << N,
|
||||
el[6] << N,
|
||||
el[7] << N,
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) fn op_or(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0] | right[0],
|
||||
left[1] | right[1],
|
||||
left[2] | right[2],
|
||||
left[3] | right[3],
|
||||
left[4] | right[4],
|
||||
left[5] | right[5],
|
||||
left[6] | right[6],
|
||||
left[7] | right[7],
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) fn op_and(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0] & right[0],
|
||||
left[1] & right[1],
|
||||
left[2] & right[2],
|
||||
left[3] & right[3],
|
||||
left[4] & right[4],
|
||||
left[5] & right[5],
|
||||
left[6] & right[6],
|
||||
left[7] & right[7],
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) unsafe fn load_unaligned(addr: *const DataType) -> DataType {
|
||||
ptr::read_unaligned(addr)
|
||||
}
|
||||
|
||||
pub(crate) unsafe fn store_unaligned(addr: *mut DataType, data: DataType) {
|
||||
ptr::write_unaligned(addr, data);
|
||||
}
|
||||
|
||||
pub(crate) fn or_collapse_to_u32(accumulator: DataType) -> u32 {
|
||||
((accumulator[0] | accumulator[1]) | (accumulator[2] | accumulator[3]))
|
||||
| ((accumulator[4] | accumulator[5]) | (accumulator[6] | accumulator[7]))
|
||||
}
|
||||
|
||||
fn compute_delta(curr: DataType, prev: DataType) -> DataType {
|
||||
[
|
||||
curr[0].wrapping_sub(prev[7]),
|
||||
curr[1].wrapping_sub(curr[0]),
|
||||
curr[2].wrapping_sub(curr[1]),
|
||||
curr[3].wrapping_sub(curr[2]),
|
||||
curr[4].wrapping_sub(curr[3]),
|
||||
curr[5].wrapping_sub(curr[4]),
|
||||
curr[6].wrapping_sub(curr[5]),
|
||||
curr[7].wrapping_sub(curr[6]),
|
||||
]
|
||||
}
|
||||
|
||||
fn integrate_delta(offset: DataType, delta: DataType) -> DataType {
|
||||
let el0 = offset[7].wrapping_add(delta[0]);
|
||||
let el1 = el0.wrapping_add(delta[1]);
|
||||
let el2 = el1.wrapping_add(delta[2]);
|
||||
let el3 = el2.wrapping_add(delta[3]);
|
||||
let el4 = el3.wrapping_add(delta[4]);
|
||||
let el5 = el4.wrapping_add(delta[5]);
|
||||
let el6 = el5.wrapping_add(delta[6]);
|
||||
let el7 = el6.wrapping_add(delta[7]);
|
||||
[el0, el1, el2, el3, el4, el5, el6, el7]
|
||||
}
|
||||
|
||||
pub(crate) fn add(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0].wrapping_add(right[0]),
|
||||
left[1].wrapping_add(right[1]),
|
||||
left[2].wrapping_add(right[2]),
|
||||
left[3].wrapping_add(right[3]),
|
||||
left[4].wrapping_add(right[4]),
|
||||
left[5].wrapping_add(right[5]),
|
||||
left[6].wrapping_add(right[6]),
|
||||
left[7].wrapping_add(right[7]),
|
||||
]
|
||||
}
|
||||
|
||||
pub(crate) fn sub(left: DataType, right: DataType) -> DataType {
|
||||
[
|
||||
left[0].wrapping_sub(right[0]),
|
||||
left[1].wrapping_sub(right[1]),
|
||||
left[2].wrapping_sub(right[2]),
|
||||
left[3].wrapping_sub(right[3]),
|
||||
left[4].wrapping_sub(right[4]),
|
||||
left[5].wrapping_sub(right[5]),
|
||||
left[6].wrapping_sub(right[6]),
|
||||
left[7].wrapping_sub(right[7]),
|
||||
]
|
||||
}
|
||||
|
||||
// The `allow(unused)` is here to put an attribute that has no effect.
|
||||
//
|
||||
// For other bitpackers, we enable a specific CPU instruction set, but for
|
||||
// the scalar bitpacker none is required.
|
||||
declare_bitpacker!(allow(unused));
|
||||
|
||||
impl Available for UnsafeBitPackerImpl {
|
||||
fn available() -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum InstructionSet {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
AVX2,
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
NEON,
|
||||
Scalar,
|
||||
}
|
||||
|
||||
/// 8-wide bitpacker implementation.
|
||||
///
|
||||
/// One block contains 256 integers.
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct BitPacker8x(InstructionSet);
|
||||
|
||||
impl BitPacker8x {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
pub(crate) fn new_avx2() -> Option<Self> {
|
||||
avx2::UnsafeBitPackerImpl::available().then_some(BitPacker8x(InstructionSet::AVX2))
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "x86_64"))]
|
||||
pub(crate) fn new_avx2() -> Option<Self> {
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
pub(crate) fn new_neon() -> Option<Self> {
|
||||
neon::UnsafeBitPackerImpl::available().then_some(BitPacker8x(InstructionSet::NEON))
|
||||
}
|
||||
|
||||
#[cfg(not(all(target_arch = "aarch64", target_endian = "little")))]
|
||||
pub(crate) fn new_neon() -> Option<Self> {
|
||||
None
|
||||
}
|
||||
|
||||
pub(crate) fn new_scalar() -> Self {
|
||||
BitPacker8x(InstructionSet::Scalar)
|
||||
}
|
||||
}
|
||||
|
||||
impl BitPacker for BitPacker8x {
|
||||
const BLOCK_LEN: usize = BLOCK_LEN;
|
||||
|
||||
fn new() -> Self {
|
||||
Self::new_avx2()
|
||||
.or_else(Self::new_neon)
|
||||
.unwrap_or_else(Self::new_scalar)
|
||||
}
|
||||
|
||||
fn compress(&self, decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => {
|
||||
avx2::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::compress(decompressed, compressed, num_bits)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn compress_sorted(
|
||||
&self,
|
||||
initial: u32,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::compress_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn compress_strictly_sorted(
|
||||
&self,
|
||||
initial: Option<u32>,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::compress_strictly_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::compress_strictly_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::compress_strictly_sorted(
|
||||
initial,
|
||||
decompressed,
|
||||
compressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decompress(&self, compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => {
|
||||
avx2::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::decompress(compressed, decompressed, num_bits)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decompress_sorted(
|
||||
&self,
|
||||
initial: u32,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::decompress_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decompress_strictly_sorted(
|
||||
&self,
|
||||
initial: Option<u32>,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::decompress_strictly_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::decompress_strictly_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::decompress_strictly_sorted(
|
||||
initial,
|
||||
compressed,
|
||||
decompressed,
|
||||
num_bits,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn num_bits(&self, decompressed: &[u32]) -> u8 {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => avx2::UnsafeBitPackerImpl::num_bits(decompressed),
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => neon::UnsafeBitPackerImpl::num_bits(decompressed),
|
||||
InstructionSet::Scalar => scalar::UnsafeBitPackerImpl::num_bits(decompressed),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn num_bits_sorted(&self, initial: u32, decompressed: &[u32]) -> u8 {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => {
|
||||
avx2::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::num_bits_sorted(initial, decompressed)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn num_bits_strictly_sorted(&self, initial: Option<u32>, decompressed: &[u32]) -> u8 {
|
||||
unsafe {
|
||||
match self.0 {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
InstructionSet::AVX2 => {
|
||||
avx2::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed)
|
||||
}
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
InstructionSet::NEON => {
|
||||
neon::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed)
|
||||
}
|
||||
InstructionSet::Scalar => {
|
||||
scalar::UnsafeBitPackerImpl::num_bits_strictly_sorted(initial, decompressed)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::BitPacker8x;
|
||||
use crate::bitpacker_internal::BitPacker;
|
||||
use bitpacking::{BitPacker as ExternalBitPacker, BitPacker8x as ExternalBitPacker8x};
|
||||
|
||||
fn mask_for_width(width: u8) -> u32 {
|
||||
match width {
|
||||
0 => 0,
|
||||
32 => u32::MAX,
|
||||
_ => (1u32 << width) - 1,
|
||||
}
|
||||
}
|
||||
|
||||
fn raw_values(width: u8, seed: u64) -> Vec<u32> {
|
||||
let mask = mask_for_width(width);
|
||||
let mut state = seed;
|
||||
(0..BitPacker8x::BLOCK_LEN)
|
||||
.map(|idx| {
|
||||
state ^= state << 13;
|
||||
state ^= state >> 7;
|
||||
state ^= state << 17;
|
||||
match seed % 4 {
|
||||
0 => 0,
|
||||
1 => mask,
|
||||
2 => idx as u32 & mask,
|
||||
_ => state as u32 & mask,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sorted_values(width: u8, seed: u64) -> (u32, Vec<u32>) {
|
||||
if width == 0 {
|
||||
return (17, vec![17; BitPacker8x::BLOCK_LEN]);
|
||||
}
|
||||
if width == 32 {
|
||||
return (0, vec![u32::MAX; BitPacker8x::BLOCK_LEN]);
|
||||
}
|
||||
|
||||
let mask = mask_for_width(width).min(127);
|
||||
let mut state = seed;
|
||||
let mut current = 17u32;
|
||||
let values = (0..BitPacker8x::BLOCK_LEN)
|
||||
.map(|_| {
|
||||
state ^= state << 13;
|
||||
state ^= state >> 7;
|
||||
state ^= state << 17;
|
||||
current += state as u32 & mask;
|
||||
current
|
||||
})
|
||||
.collect();
|
||||
(17, values)
|
||||
}
|
||||
|
||||
fn strictly_sorted_values(width: u8, seed: u64) -> (Option<u32>, Vec<u32>) {
|
||||
let mask = mask_for_width(width).min(127);
|
||||
let mut state = seed;
|
||||
let mut current = 0u32;
|
||||
let values = (0..BitPacker8x::BLOCK_LEN)
|
||||
.map(|idx| {
|
||||
if idx == 0 {
|
||||
current = 0;
|
||||
} else {
|
||||
state ^= state << 13;
|
||||
state ^= state >> 7;
|
||||
state ^= state << 17;
|
||||
current += 1 + (state as u32 & mask);
|
||||
}
|
||||
current
|
||||
})
|
||||
.collect();
|
||||
(None, values)
|
||||
}
|
||||
|
||||
fn assert_raw_compatible(ours: BitPacker8x, external: ExternalBitPacker8x) {
|
||||
for width in 0..=32 {
|
||||
for seed in [0, 1, 2, 123456789] {
|
||||
let values = raw_values(width, seed);
|
||||
assert_eq!(ours.num_bits(&values), external.num_bits(&values));
|
||||
|
||||
let mut actual = vec![0u8; BitPacker8x::compressed_block_size(width)];
|
||||
let actual_len = ours.compress(&values, &mut actual, width);
|
||||
|
||||
let mut expected = vec![0u8; ExternalBitPacker8x::compressed_block_size(width)];
|
||||
let expected_len = external.compress(&values, &mut expected, width);
|
||||
|
||||
assert_eq!(actual_len, expected_len);
|
||||
assert_eq!(actual, expected, "raw width {width} seed {seed}");
|
||||
|
||||
let mut decoded = vec![0u32; BitPacker8x::BLOCK_LEN];
|
||||
assert_eq!(ours.decompress(&actual, &mut decoded, width), actual_len);
|
||||
assert_eq!(decoded, values);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_sorted_compatible(ours: BitPacker8x, external: ExternalBitPacker8x) {
|
||||
for width in 0..=32 {
|
||||
for seed in [0, 1, 2, 123456789] {
|
||||
let (initial, values) = sorted_values(width, seed);
|
||||
assert_eq!(
|
||||
ours.num_bits_sorted(initial, &values),
|
||||
external.num_bits_sorted(initial, &values)
|
||||
);
|
||||
|
||||
let mut actual = vec![0u8; BitPacker8x::compressed_block_size(width)];
|
||||
let actual_len = ours.compress_sorted(initial, &values, &mut actual, width);
|
||||
|
||||
let mut expected = vec![0u8; ExternalBitPacker8x::compressed_block_size(width)];
|
||||
let expected_len = external.compress_sorted(initial, &values, &mut expected, width);
|
||||
|
||||
assert_eq!(actual_len, expected_len);
|
||||
assert_eq!(actual, expected, "sorted width {width} seed {seed}");
|
||||
|
||||
let mut decoded = vec![0u32; BitPacker8x::BLOCK_LEN];
|
||||
assert_eq!(
|
||||
ours.decompress_sorted(initial, &actual, &mut decoded, width),
|
||||
actual_len
|
||||
);
|
||||
assert_eq!(decoded, values);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_strictly_sorted_compatible(ours: BitPacker8x, external: ExternalBitPacker8x) {
|
||||
for width in 0..=16 {
|
||||
for seed in [0, 1, 2, 123456789] {
|
||||
let (initial, values) = strictly_sorted_values(width, seed);
|
||||
let num_bits = external.num_bits_strictly_sorted(initial, &values);
|
||||
assert_eq!(ours.num_bits_strictly_sorted(initial, &values), num_bits);
|
||||
|
||||
let mut actual = vec![0u8; BitPacker8x::compressed_block_size(num_bits)];
|
||||
let actual_len =
|
||||
ours.compress_strictly_sorted(initial, &values, &mut actual, num_bits);
|
||||
|
||||
let mut expected = vec![0u8; ExternalBitPacker8x::compressed_block_size(num_bits)];
|
||||
let expected_len =
|
||||
external.compress_strictly_sorted(initial, &values, &mut expected, num_bits);
|
||||
|
||||
assert_eq!(actual_len, expected_len);
|
||||
assert_eq!(actual, expected, "strict width {width} seed {seed}");
|
||||
|
||||
let mut decoded = vec![0u32; BitPacker8x::BLOCK_LEN];
|
||||
assert_eq!(
|
||||
ours.decompress_strictly_sorted(initial, &actual, &mut decoded, num_bits),
|
||||
actual_len
|
||||
);
|
||||
assert_eq!(decoded, values);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bitpacker8x_raw_compatible_with_external_bitpacking() {
|
||||
assert_raw_compatible(BitPacker8x::new(), ExternalBitPacker8x::new());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bitpacker8x_sorted_compatible_with_external_bitpacking() {
|
||||
assert_sorted_compatible(BitPacker8x::new(), ExternalBitPacker8x::new());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scalar_backend_matches_external_bitpacker8x() {
|
||||
let scalar = BitPacker8x::new_scalar();
|
||||
let external = ExternalBitPacker8x::new();
|
||||
|
||||
assert_raw_compatible(scalar, external);
|
||||
assert_sorted_compatible(scalar, external);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scalar_backend_matches_external_strictly_sorted_bitpacker8x() {
|
||||
let scalar = BitPacker8x::new_scalar();
|
||||
let external = ExternalBitPacker8x::new();
|
||||
|
||||
assert_strictly_sorted_compatible(scalar, external);
|
||||
}
|
||||
|
||||
#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
|
||||
#[test]
|
||||
fn neon_backend_matches_external_bitpacker8x() {
|
||||
if let Some(neon) = BitPacker8x::new_neon() {
|
||||
let external = ExternalBitPacker8x::new();
|
||||
|
||||
assert_raw_compatible(neon, external);
|
||||
assert_sorted_compatible(neon, external);
|
||||
assert_strictly_sorted_compatible(neon, external);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,672 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
macro_rules! pack_unpack_with_bits {
|
||||
|
||||
($name:ident, $n:expr, $cpufeature:meta) => {
|
||||
|
||||
|
||||
mod $name {
|
||||
|
||||
use crunchy::unroll;
|
||||
use super::BLOCK_LEN;
|
||||
use super::{Sink, Transformer};
|
||||
use super::{DataType,
|
||||
set1,
|
||||
right_shift_32,
|
||||
left_shift_32,
|
||||
op_or,
|
||||
op_and,
|
||||
load_unaligned,
|
||||
store_unaligned};
|
||||
|
||||
const NUM_BITS: usize = $n;
|
||||
const NUM_BYTES_PER_BLOCK: usize = NUM_BITS * BLOCK_LEN / 8;
|
||||
|
||||
#[$cpufeature]
|
||||
pub(crate) unsafe fn pack<TDeltaComputer: Transformer>(input_arr: &[u32], output_arr: &mut [u8], mut delta_computer: TDeltaComputer) -> usize {
|
||||
assert_eq!(input_arr.len(), BLOCK_LEN, "Input block too small {}, (expected {})", input_arr.len(), BLOCK_LEN);
|
||||
assert!(output_arr.len() >= NUM_BYTES_PER_BLOCK, "Output array too small (numbits {}). {} <= {}", NUM_BITS, output_arr.len(), NUM_BYTES_PER_BLOCK);
|
||||
|
||||
let input_ptr = input_arr.as_ptr().cast::<DataType>();
|
||||
let mut output_ptr = output_arr.as_mut_ptr().cast::<DataType>();
|
||||
let mut out_register: DataType = delta_computer.transform(load_unaligned(input_ptr));
|
||||
|
||||
unroll! {
|
||||
for iter in 0..30 {
|
||||
const i: usize = 1 + iter;
|
||||
|
||||
const bits_filled: usize = i * NUM_BITS;
|
||||
const inner_cursor: usize = bits_filled % 32;
|
||||
const remaining: usize = 32 - inner_cursor;
|
||||
|
||||
let offset_ptr = input_ptr.add(i);
|
||||
let in_register: DataType = delta_computer.transform(load_unaligned(offset_ptr));
|
||||
|
||||
out_register =
|
||||
if inner_cursor > 0 {
|
||||
op_or(out_register, left_shift_32::<{inner_cursor as i32}>(in_register))
|
||||
} else {
|
||||
in_register
|
||||
};
|
||||
|
||||
if remaining <= NUM_BITS {
|
||||
store_unaligned(output_ptr, out_register);
|
||||
output_ptr = output_ptr.add(1);
|
||||
if 0 < remaining && remaining < NUM_BITS {
|
||||
out_register = right_shift_32::<{remaining as i32}>(in_register);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let in_register: DataType = delta_computer.transform(load_unaligned(input_ptr.add(31)));
|
||||
out_register = if 32 - NUM_BITS > 0 {
|
||||
op_or(out_register, left_shift_32::<{32 - NUM_BITS as i32}>(in_register))
|
||||
} else {
|
||||
op_or(out_register, in_register)
|
||||
};
|
||||
store_unaligned(output_ptr, out_register);
|
||||
|
||||
NUM_BYTES_PER_BLOCK
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
pub(crate) unsafe fn unpack<Output: Sink>(compressed: &[u8], mut output: Output) -> usize {
|
||||
|
||||
assert!(compressed.len() >= NUM_BYTES_PER_BLOCK, "Compressed array seems too small. ({} < {}) ", compressed.len(), NUM_BYTES_PER_BLOCK);
|
||||
|
||||
let mut input_ptr = compressed.as_ptr().cast::<DataType>();
|
||||
|
||||
let mask_scalar: u32 = ((1u64 << NUM_BITS) - 1u64) as u32;
|
||||
let mask = set1(mask_scalar as i32);
|
||||
|
||||
let mut in_register: DataType = load_unaligned(input_ptr);
|
||||
|
||||
let out_register = op_and(in_register, mask);
|
||||
output.process(out_register);
|
||||
|
||||
unroll! {
|
||||
for iter in 0..31 {
|
||||
const i: usize = iter + 1;
|
||||
|
||||
const inner_cursor: usize = (i * NUM_BITS) % 32;
|
||||
const inner_capacity: usize = 32 - inner_cursor;
|
||||
|
||||
let shifted_in_register = if inner_cursor != 0 {
|
||||
right_shift_32::<{inner_cursor as i32}>(in_register)
|
||||
} else {
|
||||
in_register
|
||||
};
|
||||
let mut out_register: DataType = op_and(shifted_in_register, mask);
|
||||
|
||||
// We consumed our current quadruplets entirely.
|
||||
// We therefore read another one.
|
||||
if inner_capacity <= NUM_BITS && i != 31 {
|
||||
input_ptr = input_ptr.add(1);
|
||||
in_register = load_unaligned(input_ptr);
|
||||
|
||||
// This quadruplets is actually cutting one of
|
||||
// our `DataType`. We need to read the next one.
|
||||
if inner_capacity < NUM_BITS {
|
||||
let shifted = if inner_capacity != 0 {
|
||||
left_shift_32::<{inner_capacity as i32}>(in_register)
|
||||
} else {
|
||||
in_register
|
||||
};
|
||||
let masked = op_and(shifted, mask);
|
||||
out_register = op_or(out_register, masked);
|
||||
}
|
||||
}
|
||||
|
||||
output.process(out_register);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
NUM_BYTES_PER_BLOCK
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! pack_unpack_with_bits_32 {
|
||||
($cpufeature:meta) => {
|
||||
mod pack_unpack_with_bits_32 {
|
||||
use super::BLOCK_LEN;
|
||||
use super::{DataType, load_unaligned, store_unaligned};
|
||||
use super::{Sink, Transformer};
|
||||
use crunchy::unroll;
|
||||
|
||||
const NUM_BITS: usize = 32;
|
||||
const NUM_BYTES_PER_BLOCK: usize = NUM_BITS * BLOCK_LEN / 8;
|
||||
|
||||
#[$cpufeature]
|
||||
pub(crate) unsafe fn pack<TDeltaComputer: Transformer>(
|
||||
input_arr: &[u32],
|
||||
output_arr: &mut [u8],
|
||||
mut delta_computer: TDeltaComputer,
|
||||
) -> usize {
|
||||
assert_eq!(
|
||||
input_arr.len(),
|
||||
BLOCK_LEN,
|
||||
"Input block too small {}, (expected {})",
|
||||
input_arr.len(),
|
||||
BLOCK_LEN
|
||||
);
|
||||
assert!(
|
||||
output_arr.len() >= NUM_BYTES_PER_BLOCK,
|
||||
"Output array too small (numbits {}). {} <= {}",
|
||||
NUM_BITS,
|
||||
output_arr.len(),
|
||||
NUM_BYTES_PER_BLOCK
|
||||
);
|
||||
|
||||
let input_ptr: *const DataType = input_arr.as_ptr().cast::<DataType>();
|
||||
let output_ptr = output_arr.as_mut_ptr().cast::<DataType>();
|
||||
unroll! {
|
||||
for i in 0..32 {
|
||||
let input_offset_ptr = input_ptr.add(i);
|
||||
let output_offset_ptr = output_ptr.add(i);
|
||||
let input_register = load_unaligned(input_offset_ptr);
|
||||
let output_register = delta_computer.transform(input_register);
|
||||
store_unaligned(output_offset_ptr, output_register);
|
||||
}
|
||||
}
|
||||
NUM_BYTES_PER_BLOCK
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
pub(crate) unsafe fn unpack<Output: Sink>(
|
||||
compressed: &[u8],
|
||||
mut output: Output,
|
||||
) -> usize {
|
||||
assert!(
|
||||
compressed.len() >= NUM_BYTES_PER_BLOCK,
|
||||
"Compressed array seems too small. ({} < {}) ",
|
||||
compressed.len(),
|
||||
NUM_BYTES_PER_BLOCK
|
||||
);
|
||||
let input_ptr = compressed.as_ptr().cast::<DataType>();
|
||||
for i in 0..32 {
|
||||
let input_offset_ptr = input_ptr.add(i);
|
||||
let in_register: DataType = load_unaligned(input_offset_ptr);
|
||||
output.process(in_register);
|
||||
}
|
||||
NUM_BYTES_PER_BLOCK
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! declare_bitpacker {
|
||||
($cpufeature:meta) => {
|
||||
use super::super::UnsafeBitPacker;
|
||||
use super::super::most_significant_bit;
|
||||
use crunchy::unroll;
|
||||
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_1, 1, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_2, 2, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_3, 3, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_4, 4, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_5, 5, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_6, 6, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_7, 7, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_8, 8, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_9, 9, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_10, 10, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_11, 11, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_12, 12, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_13, 13, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_14, 14, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_15, 15, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_16, 16, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_17, 17, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_18, 18, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_19, 19, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_20, 20, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_21, 21, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_22, 22, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_23, 23, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_24, 24, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_25, 25, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_26, 26, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_27, 27, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_28, 28, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_29, 29, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_30, 30, $cpufeature);
|
||||
pack_unpack_with_bits!(pack_unpack_with_bits_31, 31, $cpufeature);
|
||||
pack_unpack_with_bits_32!($cpufeature);
|
||||
|
||||
unsafe fn compress_generic<DeltaComputer: Transformer>(
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
delta_computer: DeltaComputer,
|
||||
) -> usize {
|
||||
match num_bits {
|
||||
0 => 0,
|
||||
1 => pack_unpack_with_bits_1::pack(decompressed, compressed, delta_computer),
|
||||
2 => pack_unpack_with_bits_2::pack(decompressed, compressed, delta_computer),
|
||||
3 => pack_unpack_with_bits_3::pack(decompressed, compressed, delta_computer),
|
||||
4 => pack_unpack_with_bits_4::pack(decompressed, compressed, delta_computer),
|
||||
5 => pack_unpack_with_bits_5::pack(decompressed, compressed, delta_computer),
|
||||
6 => pack_unpack_with_bits_6::pack(decompressed, compressed, delta_computer),
|
||||
7 => pack_unpack_with_bits_7::pack(decompressed, compressed, delta_computer),
|
||||
8 => pack_unpack_with_bits_8::pack(decompressed, compressed, delta_computer),
|
||||
9 => pack_unpack_with_bits_9::pack(decompressed, compressed, delta_computer),
|
||||
10 => pack_unpack_with_bits_10::pack(decompressed, compressed, delta_computer),
|
||||
11 => pack_unpack_with_bits_11::pack(decompressed, compressed, delta_computer),
|
||||
12 => pack_unpack_with_bits_12::pack(decompressed, compressed, delta_computer),
|
||||
13 => pack_unpack_with_bits_13::pack(decompressed, compressed, delta_computer),
|
||||
14 => pack_unpack_with_bits_14::pack(decompressed, compressed, delta_computer),
|
||||
15 => pack_unpack_with_bits_15::pack(decompressed, compressed, delta_computer),
|
||||
16 => pack_unpack_with_bits_16::pack(decompressed, compressed, delta_computer),
|
||||
17 => pack_unpack_with_bits_17::pack(decompressed, compressed, delta_computer),
|
||||
18 => pack_unpack_with_bits_18::pack(decompressed, compressed, delta_computer),
|
||||
19 => pack_unpack_with_bits_19::pack(decompressed, compressed, delta_computer),
|
||||
20 => pack_unpack_with_bits_20::pack(decompressed, compressed, delta_computer),
|
||||
21 => pack_unpack_with_bits_21::pack(decompressed, compressed, delta_computer),
|
||||
22 => pack_unpack_with_bits_22::pack(decompressed, compressed, delta_computer),
|
||||
23 => pack_unpack_with_bits_23::pack(decompressed, compressed, delta_computer),
|
||||
24 => pack_unpack_with_bits_24::pack(decompressed, compressed, delta_computer),
|
||||
25 => pack_unpack_with_bits_25::pack(decompressed, compressed, delta_computer),
|
||||
26 => pack_unpack_with_bits_26::pack(decompressed, compressed, delta_computer),
|
||||
27 => pack_unpack_with_bits_27::pack(decompressed, compressed, delta_computer),
|
||||
28 => pack_unpack_with_bits_28::pack(decompressed, compressed, delta_computer),
|
||||
29 => pack_unpack_with_bits_29::pack(decompressed, compressed, delta_computer),
|
||||
30 => pack_unpack_with_bits_30::pack(decompressed, compressed, delta_computer),
|
||||
31 => pack_unpack_with_bits_31::pack(decompressed, compressed, delta_computer),
|
||||
32 => pack_unpack_with_bits_32::pack(decompressed, compressed, delta_computer),
|
||||
_ => {
|
||||
panic!("Num bits must be <= 32. Was {}.", num_bits);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait Transformer {
|
||||
unsafe fn transform(&mut self, data: DataType) -> DataType;
|
||||
}
|
||||
|
||||
struct NoDelta;
|
||||
|
||||
impl Transformer for NoDelta {
|
||||
#[inline]
|
||||
unsafe fn transform(&mut self, current: DataType) -> DataType {
|
||||
current
|
||||
}
|
||||
}
|
||||
|
||||
struct DeltaComputer {
|
||||
pub previous: DataType,
|
||||
}
|
||||
|
||||
impl Transformer for DeltaComputer {
|
||||
#[inline]
|
||||
unsafe fn transform(&mut self, current: DataType) -> DataType {
|
||||
let result = compute_delta(current, self.previous);
|
||||
self.previous = current;
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
struct StrictDeltaComputer {
|
||||
pub previous: DataType,
|
||||
}
|
||||
|
||||
impl Transformer for StrictDeltaComputer {
|
||||
#[inline]
|
||||
unsafe fn transform(&mut self, current: DataType) -> DataType {
|
||||
let result = compute_delta(current, self.previous);
|
||||
self.previous = current;
|
||||
sub(result, set1(1))
|
||||
}
|
||||
}
|
||||
|
||||
pub trait Sink {
|
||||
unsafe fn process(&mut self, data_type: DataType);
|
||||
}
|
||||
|
||||
struct Store {
|
||||
output_ptr: *mut DataType,
|
||||
}
|
||||
|
||||
impl Store {
|
||||
fn new(output_ptr: *mut DataType) -> Store {
|
||||
Store { output_ptr }
|
||||
}
|
||||
}
|
||||
|
||||
struct DeltaIntegrate {
|
||||
current: DataType,
|
||||
output_ptr: *mut DataType,
|
||||
}
|
||||
|
||||
impl DeltaIntegrate {
|
||||
unsafe fn new(initial: u32, output_ptr: *mut DataType) -> DeltaIntegrate {
|
||||
DeltaIntegrate {
|
||||
current: set1(initial as i32),
|
||||
output_ptr,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Sink for DeltaIntegrate {
|
||||
#[inline]
|
||||
unsafe fn process(&mut self, delta: DataType) {
|
||||
self.current = integrate_delta(self.current, delta);
|
||||
store_unaligned(self.output_ptr, self.current);
|
||||
self.output_ptr = self.output_ptr.add(1);
|
||||
}
|
||||
}
|
||||
|
||||
struct StrictDeltaIntegrate {
|
||||
current: DataType,
|
||||
output_ptr: *mut DataType,
|
||||
}
|
||||
|
||||
impl StrictDeltaIntegrate {
|
||||
unsafe fn new(initial: u32, output_ptr: *mut DataType) -> StrictDeltaIntegrate {
|
||||
StrictDeltaIntegrate {
|
||||
current: set1(initial as i32),
|
||||
output_ptr,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Sink for StrictDeltaIntegrate {
|
||||
#[inline]
|
||||
unsafe fn process(&mut self, delta: DataType) {
|
||||
self.current = integrate_delta(self.current, add(delta, set1(1)));
|
||||
store_unaligned(self.output_ptr, self.current);
|
||||
self.output_ptr = self.output_ptr.add(1);
|
||||
}
|
||||
}
|
||||
|
||||
impl Sink for Store {
|
||||
#[inline]
|
||||
unsafe fn process(&mut self, out_register: DataType) {
|
||||
store_unaligned(self.output_ptr, out_register);
|
||||
self.output_ptr = self.output_ptr.add(1);
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn decompress_to<Output: Sink>(
|
||||
compressed: &[u8],
|
||||
mut sink: Output,
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
match num_bits {
|
||||
0 => {
|
||||
let zero = set1(0i32);
|
||||
for _ in 0..32 {
|
||||
sink.process(zero);
|
||||
}
|
||||
0
|
||||
}
|
||||
1 => pack_unpack_with_bits_1::unpack(compressed, sink),
|
||||
2 => pack_unpack_with_bits_2::unpack(compressed, sink),
|
||||
3 => pack_unpack_with_bits_3::unpack(compressed, sink),
|
||||
4 => pack_unpack_with_bits_4::unpack(compressed, sink),
|
||||
5 => pack_unpack_with_bits_5::unpack(compressed, sink),
|
||||
6 => pack_unpack_with_bits_6::unpack(compressed, sink),
|
||||
7 => pack_unpack_with_bits_7::unpack(compressed, sink),
|
||||
8 => pack_unpack_with_bits_8::unpack(compressed, sink),
|
||||
9 => pack_unpack_with_bits_9::unpack(compressed, sink),
|
||||
10 => pack_unpack_with_bits_10::unpack(compressed, sink),
|
||||
11 => pack_unpack_with_bits_11::unpack(compressed, sink),
|
||||
12 => pack_unpack_with_bits_12::unpack(compressed, sink),
|
||||
13 => pack_unpack_with_bits_13::unpack(compressed, sink),
|
||||
14 => pack_unpack_with_bits_14::unpack(compressed, sink),
|
||||
15 => pack_unpack_with_bits_15::unpack(compressed, sink),
|
||||
16 => pack_unpack_with_bits_16::unpack(compressed, sink),
|
||||
17 => pack_unpack_with_bits_17::unpack(compressed, sink),
|
||||
18 => pack_unpack_with_bits_18::unpack(compressed, sink),
|
||||
19 => pack_unpack_with_bits_19::unpack(compressed, sink),
|
||||
20 => pack_unpack_with_bits_20::unpack(compressed, sink),
|
||||
21 => pack_unpack_with_bits_21::unpack(compressed, sink),
|
||||
22 => pack_unpack_with_bits_22::unpack(compressed, sink),
|
||||
23 => pack_unpack_with_bits_23::unpack(compressed, sink),
|
||||
24 => pack_unpack_with_bits_24::unpack(compressed, sink),
|
||||
25 => pack_unpack_with_bits_25::unpack(compressed, sink),
|
||||
26 => pack_unpack_with_bits_26::unpack(compressed, sink),
|
||||
27 => pack_unpack_with_bits_27::unpack(compressed, sink),
|
||||
28 => pack_unpack_with_bits_28::unpack(compressed, sink),
|
||||
29 => pack_unpack_with_bits_29::unpack(compressed, sink),
|
||||
30 => pack_unpack_with_bits_30::unpack(compressed, sink),
|
||||
31 => pack_unpack_with_bits_31::unpack(compressed, sink),
|
||||
32 => pack_unpack_with_bits_32::unpack(compressed, sink),
|
||||
_ => {
|
||||
panic!("Num bits must be <= 32. Was {}.", num_bits);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct UnsafeBitPackerImpl;
|
||||
|
||||
impl UnsafeBitPacker for UnsafeBitPackerImpl {
|
||||
const BLOCK_LEN: usize = BLOCK_LEN;
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn compress(decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize {
|
||||
compress_generic(decompressed, compressed, num_bits, NoDelta)
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn compress_sorted(
|
||||
initial: u32,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
let delta_computer = DeltaComputer {
|
||||
previous: set1(initial as i32),
|
||||
};
|
||||
compress_generic(decompressed, compressed, num_bits, delta_computer)
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn compress_strictly_sorted(
|
||||
initial: Option<u32>,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
// to allow encoding [0, 1, 2, ..], we need to permit an initial value "lower" than
|
||||
// zero. To get a clean api, that value is None, but in practice, as we work on
|
||||
// wrapping integers, u32::MAX/-1 does the job just fine.
|
||||
let initial = initial.unwrap_or(u32::MAX);
|
||||
let delta_computer = StrictDeltaComputer {
|
||||
previous: set1(initial as i32),
|
||||
};
|
||||
compress_generic(decompressed, compressed, num_bits, delta_computer)
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn decompress(
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
assert!(
|
||||
decompressed.len() >= BLOCK_LEN,
|
||||
"The output array is not large enough : ({} >= {})",
|
||||
decompressed.len(),
|
||||
BLOCK_LEN
|
||||
);
|
||||
let output_ptr = decompressed.as_mut_ptr().cast::<DataType>();
|
||||
let output = Store::new(output_ptr);
|
||||
decompress_to(compressed, output, num_bits)
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn decompress_sorted(
|
||||
initial: u32,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
assert!(
|
||||
decompressed.len() >= BLOCK_LEN,
|
||||
"The output array is not large enough : ({} >= {})",
|
||||
decompressed.len(),
|
||||
BLOCK_LEN
|
||||
);
|
||||
let output_ptr = decompressed.as_mut_ptr().cast::<DataType>();
|
||||
let output = DeltaIntegrate::new(initial, output_ptr);
|
||||
decompress_to(compressed, output, num_bits)
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn decompress_strictly_sorted(
|
||||
initial: Option<u32>,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize {
|
||||
assert!(
|
||||
decompressed.len() >= BLOCK_LEN,
|
||||
"The output array is not large enough : ({} >= {})",
|
||||
decompressed.len(),
|
||||
BLOCK_LEN
|
||||
);
|
||||
let initial = initial.unwrap_or(u32::MAX);
|
||||
let output_ptr = decompressed.as_mut_ptr().cast::<DataType>();
|
||||
let output = StrictDeltaIntegrate::new(initial, output_ptr);
|
||||
decompress_to(compressed, output, num_bits)
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn num_bits(decompressed: &[u32]) -> u8 {
|
||||
assert_eq!(
|
||||
decompressed.len(),
|
||||
BLOCK_LEN,
|
||||
"`decompressed`'s len is not `BLOCK_LEN={}`",
|
||||
BLOCK_LEN
|
||||
);
|
||||
let data: *const DataType = decompressed.as_ptr().cast::<DataType>();
|
||||
let mut accumulator = load_unaligned(data);
|
||||
unroll! {
|
||||
for iter in 0..31 {
|
||||
let i = iter + 1;
|
||||
let newvec = load_unaligned(data.add(i));
|
||||
accumulator = op_or(accumulator, newvec);
|
||||
}
|
||||
}
|
||||
most_significant_bit(or_collapse_to_u32(accumulator))
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn num_bits_sorted(initial: u32, decompressed: &[u32]) -> u8 {
|
||||
assert_eq!(
|
||||
decompressed.len(),
|
||||
BLOCK_LEN,
|
||||
"`decompressed`'s len is not `BLOCK_LEN={}`",
|
||||
BLOCK_LEN
|
||||
);
|
||||
let initial_vec = set1(initial as i32);
|
||||
let data: *const DataType = decompressed.as_ptr().cast::<DataType>();
|
||||
|
||||
let first = load_unaligned(data);
|
||||
let mut accumulator = compute_delta(load_unaligned(data), initial_vec);
|
||||
let mut previous = first;
|
||||
|
||||
unroll! {
|
||||
for iter in 0..30 {
|
||||
let i = iter + 1;
|
||||
let current = load_unaligned(data.add(i));
|
||||
let delta = compute_delta(current, previous);
|
||||
accumulator = op_or(accumulator, delta);
|
||||
previous = current;
|
||||
}
|
||||
}
|
||||
let current = load_unaligned(data.add(31));
|
||||
let delta = compute_delta(current, previous);
|
||||
accumulator = op_or(accumulator, delta);
|
||||
most_significant_bit(or_collapse_to_u32(accumulator))
|
||||
}
|
||||
|
||||
#[$cpufeature]
|
||||
unsafe fn num_bits_strictly_sorted(initial: Option<u32>, decompressed: &[u32]) -> u8 {
|
||||
assert_eq!(
|
||||
decompressed.len(),
|
||||
BLOCK_LEN,
|
||||
"`decompressed`'s len is not `BLOCK_LEN={}`",
|
||||
BLOCK_LEN
|
||||
);
|
||||
let initial = initial.unwrap_or(u32::MAX);
|
||||
let initial_vec = set1(initial as i32);
|
||||
let one = set1(1);
|
||||
let data: *const DataType = decompressed.as_ptr().cast::<DataType>();
|
||||
|
||||
let first = load_unaligned(data);
|
||||
let mut accumulator = sub(compute_delta(load_unaligned(data), initial_vec), one);
|
||||
let mut previous = first;
|
||||
|
||||
unroll! {
|
||||
for iter in 0..30 {
|
||||
let i = iter + 1;
|
||||
let current = load_unaligned(data.add(i));
|
||||
let delta = sub(compute_delta(current, previous), one);
|
||||
accumulator = op_or(accumulator, delta);
|
||||
previous = current;
|
||||
}
|
||||
}
|
||||
let current = load_unaligned(data.add(31));
|
||||
let delta = sub(compute_delta(current, previous), one);
|
||||
accumulator = op_or(accumulator, delta);
|
||||
most_significant_bit(or_collapse_to_u32(accumulator))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any())]
|
||||
mod tests {
|
||||
use super::super::UnsafeBitPacker;
|
||||
use super::UnsafeBitPackerImpl;
|
||||
use crate::Available;
|
||||
use crate::tests::{DeltaKind, test_suite_compress_decompress};
|
||||
|
||||
#[test]
|
||||
fn test_num_bits() {
|
||||
if UnsafeBitPackerImpl::available() {
|
||||
for num_bits in 0..32 {
|
||||
for pos in 0..32 {
|
||||
let mut vals = [0u32; UnsafeBitPackerImpl::BLOCK_LEN];
|
||||
if num_bits > 0 {
|
||||
vals[pos] = 1 << (num_bits - 1);
|
||||
}
|
||||
assert_eq!(
|
||||
unsafe { UnsafeBitPackerImpl::num_bits(&vals[..]) },
|
||||
num_bits
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bitpacker() {
|
||||
if UnsafeBitPackerImpl::available() {
|
||||
test_suite_compress_decompress::<UnsafeBitPackerImpl>(DeltaKind::NoDelta);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bitpacker_delta() {
|
||||
if UnsafeBitPackerImpl::available() {
|
||||
test_suite_compress_decompress::<UnsafeBitPackerImpl>(DeltaKind::Delta);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bitpacker_strict_delta() {
|
||||
if UnsafeBitPackerImpl::available() {
|
||||
test_suite_compress_decompress::<UnsafeBitPackerImpl>(DeltaKind::StrictDelta);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
// Lance-owned u32 SIMD bitpacking kernels.
|
||||
//
|
||||
// This is adapted from the MIT-licensed `bitpacking` crate so Lance can keep the
|
||||
// hot FTS posting-list bitpacking implementation inside lance-bitpacking while
|
||||
// preserving byte compatibility with the existing 4x format.
|
||||
|
||||
#![allow(dead_code)]
|
||||
#![allow(unsafe_op_in_unsafe_fn)]
|
||||
#![allow(clippy::redundant_pub_crate)]
|
||||
#![allow(clippy::upper_case_acronyms)]
|
||||
#![allow(clippy::use_self)]
|
||||
|
||||
#[macro_use]
|
||||
mod macros;
|
||||
|
||||
mod bitpacker4x;
|
||||
mod bitpacker8x;
|
||||
|
||||
pub use bitpacker4x::BitPacker4x;
|
||||
pub use bitpacker8x::BitPacker8x;
|
||||
|
||||
pub(crate) trait Available {
|
||||
fn available() -> bool;
|
||||
}
|
||||
|
||||
pub(crate) trait UnsafeBitPacker {
|
||||
const BLOCK_LEN: usize;
|
||||
|
||||
unsafe fn compress(decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize;
|
||||
|
||||
unsafe fn compress_sorted(
|
||||
initial: u32,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
unsafe fn compress_strictly_sorted(
|
||||
initial: Option<u32>,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
unsafe fn decompress(compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize;
|
||||
|
||||
unsafe fn decompress_sorted(
|
||||
initial: u32,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
unsafe fn decompress_strictly_sorted(
|
||||
initial: Option<u32>,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
unsafe fn num_bits(decompressed: &[u32]) -> u8;
|
||||
|
||||
unsafe fn num_bits_sorted(initial: u32, decompressed: &[u32]) -> u8;
|
||||
|
||||
unsafe fn num_bits_strictly_sorted(initial: Option<u32>, decompressed: &[u32]) -> u8;
|
||||
}
|
||||
|
||||
/// Block bitpacker for fixed-size `u32` blocks.
|
||||
///
|
||||
/// Implementations own runtime SIMD dispatch and use caller-provided buffers.
|
||||
/// Packed bytes are stable for a given implementation and bit width.
|
||||
pub trait BitPacker: Sized + Clone + Copy {
|
||||
/// Number of `u32` values in one physical block.
|
||||
const BLOCK_LEN: usize;
|
||||
|
||||
/// Select the best supported implementation for the current CPU.
|
||||
///
|
||||
/// Lance uses SIMD backends when available and falls back to a scalar
|
||||
/// backend otherwise, matching the existing allocation-free call shape used
|
||||
/// by the upstream `bitpacking` crate.
|
||||
fn new() -> Self;
|
||||
|
||||
/// Compress one full block of raw values into `compressed`.
|
||||
fn compress(&self, decompressed: &[u32], compressed: &mut [u8], num_bits: u8) -> usize;
|
||||
|
||||
/// Delta-compress one full non-decreasing block into `compressed`.
|
||||
fn compress_sorted(
|
||||
&self,
|
||||
initial: u32,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
/// Delta-compress one full strictly increasing block into `compressed`.
|
||||
fn compress_strictly_sorted(
|
||||
&self,
|
||||
initial: Option<u32>,
|
||||
decompressed: &[u32],
|
||||
compressed: &mut [u8],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
/// Decompress one raw block into `decompressed`.
|
||||
fn decompress(&self, compressed: &[u8], decompressed: &mut [u32], num_bits: u8) -> usize;
|
||||
|
||||
/// Decompress one delta-compressed non-decreasing block.
|
||||
fn decompress_sorted(
|
||||
&self,
|
||||
initial: u32,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
/// Decompress one delta-compressed strictly increasing block.
|
||||
fn decompress_strictly_sorted(
|
||||
&self,
|
||||
initial: Option<u32>,
|
||||
compressed: &[u8],
|
||||
decompressed: &mut [u32],
|
||||
num_bits: u8,
|
||||
) -> usize;
|
||||
|
||||
/// Return the minimum bit width needed to represent a full raw block.
|
||||
fn num_bits(&self, decompressed: &[u32]) -> u8;
|
||||
|
||||
/// Return the minimum bit width needed to represent deltas in a full block.
|
||||
fn num_bits_sorted(&self, initial: u32, decompressed: &[u32]) -> u8;
|
||||
|
||||
/// Return the minimum bit width needed to represent strict deltas in a full block.
|
||||
fn num_bits_strictly_sorted(&self, initial: Option<u32>, decompressed: &[u32]) -> u8;
|
||||
|
||||
/// Return the byte size of one compressed block at `num_bits`.
|
||||
#[must_use]
|
||||
fn compressed_block_size(num_bits: u8) -> usize {
|
||||
Self::BLOCK_LEN * num_bits as usize / 8
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn most_significant_bit(value: u32) -> u8 {
|
||||
(u32::BITS - value.leading_zeros()) as u8
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,27 @@
|
||||
[package]
|
||||
name = "fsst"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
readme.workspace = true
|
||||
repository.workspace = true
|
||||
description = "FSST string compression for Lance"
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
rust-version.workspace = true
|
||||
|
||||
[dependencies]
|
||||
arrow-array.workspace = true
|
||||
rand.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
test-log.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[[example]]
|
||||
name = "benchmark"
|
||||
path = "examples/benchmark.rs"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,160 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use std::fs::File;
|
||||
use std::io::{BufRead, BufReader};
|
||||
|
||||
use arrow_array::StringArray;
|
||||
use fsst::fsst::{FSST_SYMBOL_TABLE_SIZE, compress, decompress};
|
||||
use rand::Rng;
|
||||
|
||||
const TEST_NUM: usize = 20;
|
||||
const BUFFER_SIZE: usize = 8 * 1024 * 1024;
|
||||
|
||||
fn read_random_8_m_chunk(file_path: &str) -> Result<StringArray, std::io::Error> {
|
||||
let file = File::open(file_path)?;
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
let lines: Vec<String> = reader.lines().collect::<std::result::Result<_, _>>()?;
|
||||
let num_lines = lines.len();
|
||||
|
||||
let mut rng = rand::rng();
|
||||
let mut curr_line = rng.random_range(0..num_lines);
|
||||
|
||||
let chunk_size = BUFFER_SIZE;
|
||||
let mut size = 0;
|
||||
let mut result_lines = vec![];
|
||||
while size + lines[curr_line].len() < chunk_size {
|
||||
result_lines.push(lines[curr_line].clone());
|
||||
size += lines[curr_line].len();
|
||||
curr_line += 1;
|
||||
curr_line %= num_lines;
|
||||
}
|
||||
|
||||
Ok(StringArray::from(result_lines))
|
||||
}
|
||||
|
||||
fn benchmark(file_path: &str) {
|
||||
// Step 1: load data in memory
|
||||
let mut inputs: Vec<StringArray> = vec![];
|
||||
let mut symbol_tables: Vec<[u8; FSST_SYMBOL_TABLE_SIZE]> = vec![];
|
||||
for _ in 0..TEST_NUM {
|
||||
let this_input = read_random_8_m_chunk(file_path).unwrap();
|
||||
inputs.push(this_input);
|
||||
symbol_tables.push([0u8; FSST_SYMBOL_TABLE_SIZE]);
|
||||
}
|
||||
|
||||
// Step 2: allocate memory for compression and decompression outputs
|
||||
let mut compression_out_bufs = vec![];
|
||||
let mut compression_out_offsets_bufs = vec![];
|
||||
for _ in 0..TEST_NUM {
|
||||
let this_com_out_buf = vec![0u8; BUFFER_SIZE];
|
||||
let this_com_out_offsets_buf = vec![0i32; BUFFER_SIZE];
|
||||
compression_out_bufs.push(this_com_out_buf);
|
||||
compression_out_offsets_bufs.push(this_com_out_offsets_buf);
|
||||
}
|
||||
let mut decompression_out_bufs = vec![];
|
||||
let mut decompression_out_offsets_bufs = vec![];
|
||||
for _ in 0..TEST_NUM {
|
||||
// `decompress` requires the output buffer to be at least 8x the compressed input (a 1-byte
|
||||
// code can expand to an 8-byte symbol). The compressed buffer is at most `BUFFER_SIZE`, so
|
||||
// `BUFFER_SIZE * 8` is a safe upper bound.
|
||||
let this_decom_out_buf = vec![0u8; BUFFER_SIZE * 8];
|
||||
let this_decom_out_offsets_buf = vec![0i32; BUFFER_SIZE * 3];
|
||||
decompression_out_bufs.push(this_decom_out_buf);
|
||||
decompression_out_offsets_bufs.push(this_decom_out_offsets_buf);
|
||||
}
|
||||
|
||||
let original_total_size: usize = inputs.iter().map(|input| input.values().len()).sum();
|
||||
|
||||
// Step 3: compress data
|
||||
let start = std::time::Instant::now();
|
||||
for i in 0..TEST_NUM {
|
||||
compress(
|
||||
symbol_tables[i].as_mut(),
|
||||
inputs[i].values(),
|
||||
inputs[i].value_offsets(),
|
||||
&mut compression_out_bufs[i],
|
||||
&mut compression_out_offsets_bufs[i],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
let compression_finish_time = std::time::Instant::now();
|
||||
|
||||
for i in 0..TEST_NUM {
|
||||
decompress(
|
||||
&symbol_tables[i],
|
||||
&compression_out_bufs[i],
|
||||
&compression_out_offsets_bufs[i],
|
||||
&mut decompression_out_bufs[i],
|
||||
&mut decompression_out_offsets_bufs[i],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
let decompression_finish_time = std::time::Instant::now();
|
||||
let compression_total_size: usize = compression_out_bufs.iter().map(|buf| buf.len()).sum();
|
||||
let compression_ratio = original_total_size as f64 / compression_total_size as f64;
|
||||
let compress_time = compression_finish_time - start;
|
||||
let decompress_time = decompression_finish_time - compression_finish_time;
|
||||
|
||||
let compress_seconds =
|
||||
compress_time.as_secs() as f64 + compress_time.subsec_nanos() as f64 * 1e-9;
|
||||
|
||||
let decompress_seconds =
|
||||
decompress_time.as_secs() as f64 + decompress_time.subsec_nanos() as f64 * 1e-9;
|
||||
|
||||
let com_speed = (original_total_size as f64 / compress_seconds) / 1024f64 / 1024f64;
|
||||
|
||||
let d_speed = (original_total_size as f64 / decompress_seconds) / 1024f64 / 1024f64;
|
||||
for i in 0..TEST_NUM {
|
||||
assert_eq!(
|
||||
inputs[i].value_offsets().len(),
|
||||
decompression_out_offsets_bufs[i].len()
|
||||
);
|
||||
}
|
||||
|
||||
// Print tsv headers
|
||||
#[allow(clippy::print_stdout)]
|
||||
{
|
||||
println!("for file: {}", file_path);
|
||||
println!("Compression ratio\tCompression speed\tDecompression speed");
|
||||
println!(
|
||||
"{:.3}\t\t\t\t{:.2}MB/s\t\t\t{:.2}MB/s",
|
||||
compression_ratio, com_speed, d_speed
|
||||
);
|
||||
}
|
||||
for i in 0..TEST_NUM {
|
||||
assert_eq!(inputs[i].value_data(), decompression_out_bufs[i]);
|
||||
assert_eq!(inputs[i].value_offsets(), decompression_out_offsets_bufs[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// to run this test, download MS Marco dataset from https://msmarco.z22.web.core.windows.net/msmarcoranking/fulldocs.tsv.gz
|
||||
// and use a script like this to get each column
|
||||
/*
|
||||
import csv
|
||||
import sys
|
||||
def write_second_column(input_path, output_path):
|
||||
csv.field_size_limit(sys.maxsize)
|
||||
with open(input_path, 'r') as input_file, open(output_path, 'w') as output_file:
|
||||
tsv_reader = csv.reader(input_file, delimiter='\t')
|
||||
tsv_writer = csv.writer(output_file, delimiter='\t')
|
||||
for row in tsv_reader:
|
||||
tsv_writer.writerow([row[2]])
|
||||
#write_second_column('/Users/x/fulldocs.tsv', '/Users/x/first_column_fulldocs.tsv')
|
||||
#write_second_column('/Users/x/fulldocs.tsv', '/Users/x/second_column_fulldocs.tsv')
|
||||
write_second_column('/Users/x/fulldocs.tsv', '/Users/x/third_column_fulldocs.tsv')
|
||||
*/
|
||||
fn main() {
|
||||
let file_paths = [
|
||||
"/home/x/first_column_fulldocs.tsv",
|
||||
"/home/x/second_column_fulldocs.tsv",
|
||||
"/home/x/third_column_fulldocs_chunk_0.tsv",
|
||||
];
|
||||
for file_path in file_paths {
|
||||
benchmark(file_path);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
pub mod fsst;
|
||||
@@ -0,0 +1,54 @@
|
||||
[package]
|
||||
name = "lance-examples"
|
||||
description = "Lance examples in Rust"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
readme.workspace = true
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
rust-version.workspace = true
|
||||
|
||||
[[example]]
|
||||
name = "full_text_search"
|
||||
path = "src/full_text_search.rs"
|
||||
|
||||
[[example]]
|
||||
name = "hnsw"
|
||||
path = "src/hnsw.rs"
|
||||
|
||||
[[example]]
|
||||
name = "ivf_hnsw"
|
||||
path = "src/ivf_hnsw.rs"
|
||||
|
||||
[[example]]
|
||||
name = "llm_dataset_creation"
|
||||
path = "src/llm_dataset_creation.rs"
|
||||
|
||||
[[example]]
|
||||
name = "write_read_ds"
|
||||
path = "src/write_read_ds.rs"
|
||||
|
||||
[dependencies]
|
||||
arrow = { workspace = true }
|
||||
arrow-schema = { workspace = true }
|
||||
arrow-select = { workspace = true }
|
||||
clap = { workspace = true, features = ["derive"] }
|
||||
itertools = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
lance = { workspace = true, features = ["aws", "azure", "gcp", "oss", "huggingface", "tencent", "goosefs"] }
|
||||
lance-index = { workspace = true }
|
||||
lance-core = { workspace = true }
|
||||
lance-linalg = { workspace = true }
|
||||
lance-datagen = { workspace = true }
|
||||
object_store = {workspace = true}
|
||||
tempfile = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
all_asserts = "2.3.1"
|
||||
env_logger = "0.11.7"
|
||||
hf-hub = "0.4.2"
|
||||
parquet = { version = "58.0.0", default-features = false, features = ["arrow", "async"] }
|
||||
tokenizers = "0.15.2"
|
||||
rand.workspace = true
|
||||
@@ -0,0 +1,132 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Benchmark of HNSW graph.
|
||||
//!
|
||||
//!
|
||||
#![allow(clippy::print_stdout)]
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use all_asserts::assert_gt;
|
||||
use arrow::array::AsArray;
|
||||
use arrow::array::{Array, LargeStringArray, RecordBatch, RecordBatchIterator, UInt64Array};
|
||||
use arrow::datatypes::UInt64Type;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use itertools::Itertools;
|
||||
use lance::Dataset;
|
||||
use lance::index::DatasetIndexExt;
|
||||
use lance_datagen::{RowCount, array};
|
||||
use lance_index::scalar::inverted::flat_full_text_search;
|
||||
use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams};
|
||||
use object_store::path::Path;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
env_logger::init();
|
||||
|
||||
const TOTAL: usize = 10_000;
|
||||
const DOC_ID: &str = "__example_doc_id";
|
||||
|
||||
let tempdir = tempfile::tempdir().unwrap();
|
||||
let dataset_dir = Path::from_filesystem_path(tempdir.path()).unwrap();
|
||||
|
||||
let create_index = true;
|
||||
if create_index {
|
||||
let row_id_col = Arc::new(UInt64Array::from(
|
||||
(0..TOTAL).map(|i| i as u64).collect_vec(),
|
||||
));
|
||||
|
||||
// Generate random words using lance-datagen
|
||||
let mut words_gen = array::random_sentence(1, 100, true);
|
||||
let doc_col = words_gen
|
||||
.generate_default(RowCount::from(TOTAL as u64))
|
||||
.unwrap();
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("doc", DataType::LargeUtf8, false),
|
||||
Field::new(DOC_ID, DataType::UInt64, false),
|
||||
])),
|
||||
vec![doc_col.clone(), row_id_col.clone()],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let batches = RecordBatchIterator::new([Ok(batch.clone())], batch.schema());
|
||||
let mut dataset = Dataset::write(batches, dataset_dir.as_ref(), None)
|
||||
.await
|
||||
.unwrap();
|
||||
let params = InvertedIndexParams::default();
|
||||
let start = std::time::Instant::now();
|
||||
dataset
|
||||
.create_index(
|
||||
&["doc"],
|
||||
lance_index::IndexType::Inverted,
|
||||
None,
|
||||
¶ms,
|
||||
true,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
println!("create_index: {:?}", start.elapsed());
|
||||
}
|
||||
|
||||
let dataset = Dataset::open(dataset_dir.as_ref()).await.unwrap();
|
||||
// Use a sample word for query - fetch first doc and pick a word from it
|
||||
let sample_batch = dataset
|
||||
.scan()
|
||||
.project(&["doc"])
|
||||
.unwrap()
|
||||
.limit(Some(1), None)
|
||||
.unwrap()
|
||||
.try_into_batch()
|
||||
.await
|
||||
.unwrap();
|
||||
let sample_doc = sample_batch["doc"]
|
||||
.as_any()
|
||||
.downcast_ref::<LargeStringArray>()
|
||||
.unwrap()
|
||||
.value(0);
|
||||
let query_string = sample_doc.split_whitespace().next().unwrap();
|
||||
let query = FullTextSearchQuery::new(query_string.to_owned()).limit(Some(10));
|
||||
println!("query: {:?}", query);
|
||||
let batch = dataset
|
||||
.scan()
|
||||
.full_text_search(query.clone())
|
||||
.unwrap()
|
||||
.try_into_batch()
|
||||
.await
|
||||
.unwrap();
|
||||
let index_results = batch[DOC_ID]
|
||||
.as_primitive::<UInt64Type>()
|
||||
.iter()
|
||||
.map(|v| v.unwrap())
|
||||
.collect::<HashSet<_>>();
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
dataset
|
||||
.scan()
|
||||
.full_text_search(query.clone())
|
||||
.unwrap()
|
||||
.try_into_batch()
|
||||
.await
|
||||
.unwrap();
|
||||
println!("full_text_search: {:?}", start.elapsed());
|
||||
|
||||
let batch = dataset
|
||||
.scan()
|
||||
.project(&["doc"])
|
||||
.unwrap()
|
||||
.with_row_id()
|
||||
.try_into_batch()
|
||||
.await
|
||||
.unwrap();
|
||||
let flat_results = flat_full_text_search(&[&batch], "doc", query_string, None)
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.collect::<HashSet<_>>();
|
||||
assert_gt!(index_results.len(), 0);
|
||||
assert_eq!(index_results, flat_results);
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! HNSW is a graph based algorithm for approximate neighbor search in high-dimensional spaces.
|
||||
//! In this example, we will demonstrate how to build HNSW vector indexing against a Lance dataset.
|
||||
//! run with `cargo run -v --package lance-examples --example hnsw``
|
||||
// linked to `docs/examples/Rust/hnsw.rst`
|
||||
#![allow(clippy::print_stdout)]
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow::array::{Array, FixedSizeListArray, types::Float32Type};
|
||||
use arrow::array::{AsArray, FixedSizeListBuilder, Float32Builder};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use arrow::record_batch::RecordBatchIterator;
|
||||
use arrow_select::concat::concat;
|
||||
use futures::stream::StreamExt;
|
||||
use lance::Dataset;
|
||||
use lance_index::vector::v3::subindex::IvfSubIndex;
|
||||
use lance_index::vector::{
|
||||
flat::storage::FlatFloatStorage,
|
||||
hnsw::{
|
||||
HNSW,
|
||||
builder::{HnswBuildParams, HnswQueryParams},
|
||||
},
|
||||
};
|
||||
use lance_linalg::distance::DistanceType;
|
||||
|
||||
fn ground_truth(fsl: &FixedSizeListArray, query: &[f32], k: usize) -> HashSet<u32> {
|
||||
let mut dists = vec![];
|
||||
for i in 0..fsl.len() {
|
||||
let dist = lance_linalg::distance::l2_distance(
|
||||
query,
|
||||
fsl.value(i).as_primitive::<Float32Type>().values(),
|
||||
);
|
||||
dists.push((dist, i as u32));
|
||||
}
|
||||
dists.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
|
||||
dists.truncate(k);
|
||||
dists.into_iter().map(|(_, i)| i).collect()
|
||||
}
|
||||
|
||||
pub async fn create_test_vector_dataset(output: &str, num_rows: usize, dim: i32) {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim),
|
||||
false,
|
||||
)]));
|
||||
|
||||
let mut batches = Vec::new();
|
||||
|
||||
// Create a few batches
|
||||
for _ in 0..2 {
|
||||
let v_builder = Float32Builder::new();
|
||||
let mut list_builder = FixedSizeListBuilder::new(v_builder, dim);
|
||||
|
||||
for _ in 0..num_rows {
|
||||
for _ in 0..dim {
|
||||
list_builder.values().append_value(rand::random::<f32>());
|
||||
}
|
||||
list_builder.append(true);
|
||||
}
|
||||
let array = Arc::new(list_builder.finish());
|
||||
let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
|
||||
batches.push(batch);
|
||||
}
|
||||
let batch_reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone());
|
||||
println!("Writing dataset to {}", output);
|
||||
Dataset::write(batch_reader, output, None).await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let uri: Option<String> = None; // None means generate test data
|
||||
let column = "vector";
|
||||
let ef = 100;
|
||||
let max_edges = 30;
|
||||
let max_level = 7;
|
||||
|
||||
// 1. Generate a synthetic test data of specified dimensions
|
||||
let dataset = match uri.as_deref() {
|
||||
None => {
|
||||
println!("No uri is provided, generating test dataset...");
|
||||
let output = "test_vectors.lance";
|
||||
create_test_vector_dataset(output, 1000, 64).await;
|
||||
Dataset::open(output).await.expect("Failed to open dataset")
|
||||
}
|
||||
Some(uri) => Dataset::open(uri).await.expect("Failed to open dataset"),
|
||||
};
|
||||
|
||||
println!("Dataset schema: {:#?}", dataset.schema());
|
||||
let batches = dataset
|
||||
.scan()
|
||||
.project(&[column])
|
||||
.unwrap()
|
||||
.try_into_stream()
|
||||
.await
|
||||
.unwrap()
|
||||
.then(|batch| async move { batch.unwrap().column_by_name(column).unwrap().clone() })
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
let arrs = batches.iter().map(|b| b.as_ref()).collect::<Vec<_>>();
|
||||
let fsl = concat(&arrs).unwrap().as_fixed_size_list().clone();
|
||||
println!("Loaded {:?} batches", fsl.len());
|
||||
|
||||
let vector_store = Arc::new(FlatFloatStorage::new(fsl.clone(), DistanceType::L2));
|
||||
|
||||
let q = fsl.value(0);
|
||||
let k = 10;
|
||||
let gt = ground_truth(&fsl, q.as_primitive::<Float32Type>().values(), k);
|
||||
|
||||
for ef_construction in [15, 30, 50] {
|
||||
let now = std::time::Instant::now();
|
||||
// 2. Build a hierarchical graph structure for efficient vector search using Lance API
|
||||
let hnsw = HNSW::index_vectors(
|
||||
vector_store.as_ref(),
|
||||
HnswBuildParams::default()
|
||||
.max_level(max_level)
|
||||
.num_edges(max_edges)
|
||||
.ef_construction(ef_construction),
|
||||
)
|
||||
.unwrap();
|
||||
let construct_time = now.elapsed().as_secs_f32();
|
||||
let now = std::time::Instant::now();
|
||||
// 3. Perform vector search with different parameters and compute the ground truth using L2 distance search
|
||||
let params = HnswQueryParams {
|
||||
ef,
|
||||
lower_bound: None,
|
||||
upper_bound: None,
|
||||
dist_q_c: 0.0,
|
||||
use_acorn: false,
|
||||
};
|
||||
let results: HashSet<u32> = hnsw
|
||||
.search_basic(q.clone(), k, ¶ms, None, vector_store.as_ref())
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|node| node.id)
|
||||
.collect();
|
||||
let search_time = now.elapsed().as_micros();
|
||||
println!(
|
||||
"level={}, ef_construct={}, ef={} recall={}: construct={:.3}s search={:.3} us",
|
||||
max_level,
|
||||
ef_construction,
|
||||
ef,
|
||||
results.intersection(>).count() as f32 / k as f32,
|
||||
construct_time,
|
||||
search_time
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Run recall benchmarks for HNSW.
|
||||
//!
|
||||
//! run with `cargo run --release --example hnsw`
|
||||
#![allow(clippy::print_stdout)]
|
||||
use arrow::array::AsArray;
|
||||
use arrow::array::types::Float32Type;
|
||||
use clap::Parser;
|
||||
use futures::TryStreamExt;
|
||||
use lance::Dataset;
|
||||
use lance::dataset::ProjectionRequest;
|
||||
use lance::index::DatasetIndexExt;
|
||||
use lance::index::vector::VectorIndexParams;
|
||||
use lance_index::IndexType;
|
||||
use lance_index::vector::hnsw::builder::HnswBuildParams;
|
||||
use lance_index::vector::ivf::IvfBuildParams;
|
||||
use lance_index::vector::sq::builder::SQBuildParams;
|
||||
use lance_linalg::distance::MetricType;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(version, about, long_about = None)]
|
||||
struct Args {
|
||||
/// Dataset URI
|
||||
uri: String,
|
||||
|
||||
/// Vector column name
|
||||
#[arg(short, long, value_name = "NAME", default_value = "vector")]
|
||||
column: Option<String>,
|
||||
|
||||
#[arg(long, default_value = "100")]
|
||||
ef: usize,
|
||||
|
||||
/// Max number of edges of each node.
|
||||
#[arg(long, default_value = "30")]
|
||||
max_edges: usize,
|
||||
|
||||
#[arg(long, default_value = "7")]
|
||||
max_level: u16,
|
||||
|
||||
#[arg(long, default_value = "1")]
|
||||
nprobe: usize,
|
||||
|
||||
#[arg(short, default_value = "10")]
|
||||
k: usize,
|
||||
|
||||
#[arg(long, default_value = "false")]
|
||||
create_index: bool,
|
||||
|
||||
#[arg(long, default_value = "cosine")]
|
||||
metric_type: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn ground_truth(mat: &MatrixView<Float32Type>, query: &[f32], k: usize) -> HashSet<u32> {
|
||||
let mut dists = vec![];
|
||||
for i in 0..mat.num_rows() {
|
||||
let dist = lance_linalg::distance::l2_distance(query, mat.row(i).unwrap());
|
||||
dists.push((dist, i as u32));
|
||||
}
|
||||
dists.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
|
||||
dists.truncate(k);
|
||||
dists.into_iter().map(|(_, i)| i).collect()
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
env_logger::init();
|
||||
let args = Args::parse();
|
||||
|
||||
let mut dataset = Dataset::open(&args.uri)
|
||||
.await
|
||||
.expect("Failed to open dataset");
|
||||
println!("Dataset schema: {:#?}", dataset.schema());
|
||||
|
||||
let column = args.column.as_deref().unwrap_or("vector");
|
||||
let metric_type = MetricType::try_from(args.metric_type.as_str()).unwrap();
|
||||
|
||||
let mut ivf_params = IvfBuildParams::new(128);
|
||||
ivf_params.sample_rate = 20480;
|
||||
let hnsw_params = HnswBuildParams::default()
|
||||
.ef_construction(100)
|
||||
.num_edges(15);
|
||||
let pq_params = SQBuildParams::default();
|
||||
let params =
|
||||
VectorIndexParams::with_ivf_hnsw_sq_params(metric_type, ivf_params, hnsw_params, pq_params);
|
||||
println!("{:?}", params);
|
||||
|
||||
if args.create_index {
|
||||
let now = std::time::Instant::now();
|
||||
dataset
|
||||
.create_index(&[column], IndexType::Vector, None, ¶ms, true)
|
||||
.await
|
||||
.unwrap();
|
||||
println!("build={:.3}s", now.elapsed().as_secs_f32());
|
||||
}
|
||||
|
||||
println!("Loaded {} records", dataset.count_rows(None).await.unwrap());
|
||||
|
||||
let take_projection = ProjectionRequest::from_columns([column], dataset.schema());
|
||||
|
||||
let q = dataset
|
||||
.take(&[0], take_projection)
|
||||
.await
|
||||
.unwrap()
|
||||
.column(0)
|
||||
.as_fixed_size_list()
|
||||
.values()
|
||||
.as_primitive::<Float32Type>()
|
||||
.clone();
|
||||
|
||||
let columns: &[&str] = &[];
|
||||
let mut scan = dataset.scan();
|
||||
let plan = scan
|
||||
.project(columns)
|
||||
.unwrap()
|
||||
.with_row_id()
|
||||
.nearest(column, &q, args.k)
|
||||
.unwrap()
|
||||
.minimum_nprobes(args.nprobe);
|
||||
println!("{:?}", plan.explain_plan(true).await.unwrap());
|
||||
|
||||
let now = std::time::Instant::now();
|
||||
plan.try_into_stream()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
println!(
|
||||
"level={}, nprobe={}, k={}, search={:?}",
|
||||
args.max_level,
|
||||
args.nprobe,
|
||||
args.k,
|
||||
now.elapsed(),
|
||||
);
|
||||
|
||||
let now = std::time::Instant::now();
|
||||
for _ in 0..10 {
|
||||
plan.try_into_stream()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
println!(
|
||||
"warm up: level={}, nprobe={}, k={}, search={:?}",
|
||||
args.max_level,
|
||||
args.nprobe,
|
||||
args.k,
|
||||
now.elapsed().div_f32(10.0),
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,259 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! This example demonstrates how to:
|
||||
//!
|
||||
//! 1. Download and process a text dataset in parts from huggingface
|
||||
//! 2. Tokenize the text data with a custom RecordBatchReader
|
||||
//! 3. Save it as a Lance dataset using Lance API
|
||||
//!
|
||||
//! Run with `cargo run -v --package lance-examples --example llm_dataset_creation`
|
||||
//!
|
||||
//!
|
||||
// linked to `docs/examples/Rust/llm_dataset_creation.rst`
|
||||
|
||||
use arrow::array::{Array, Int64Builder, ListBuilder, UInt32Array};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use arrow::record_batch::RecordBatch;
|
||||
use arrow::record_batch::RecordBatchReader;
|
||||
use futures::StreamExt;
|
||||
use hf_hub::{Repo, RepoType, api::sync::Api};
|
||||
use lance::dataset::WriteParams;
|
||||
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
||||
use rand::SeedableRng;
|
||||
use rand::seq::SliceRandom;
|
||||
use std::error::Error;
|
||||
use std::fs::File;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
use tempfile::NamedTempFile;
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
// Implement a custom stream batch reader
|
||||
struct WikiTextBatchReader {
|
||||
schema: Arc<Schema>,
|
||||
parquet_readers: Vec<Option<ParquetRecordBatchReaderBuilder<File>>>,
|
||||
current_reader_idx: usize,
|
||||
current_reader: Option<Box<dyn RecordBatchReader + Send>>,
|
||||
tokenizer: Tokenizer,
|
||||
num_samples: u64,
|
||||
cur_samples_cnt: u64,
|
||||
}
|
||||
|
||||
impl WikiTextBatchReader {
|
||||
fn new(
|
||||
parquet_readers: Vec<ParquetRecordBatchReaderBuilder<File>>,
|
||||
tokenizer: Tokenizer,
|
||||
num_samples: Option<u64>,
|
||||
) -> Result<Self, Box<dyn Error + Send + Sync>> {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"input_ids",
|
||||
DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
|
||||
false,
|
||||
)]));
|
||||
|
||||
Ok(Self {
|
||||
schema,
|
||||
parquet_readers: parquet_readers.into_iter().map(Some).collect(),
|
||||
current_reader_idx: 0,
|
||||
current_reader: None,
|
||||
tokenizer,
|
||||
num_samples: num_samples.unwrap_or(100_000),
|
||||
cur_samples_cnt: 0,
|
||||
})
|
||||
}
|
||||
|
||||
fn process_batch(
|
||||
&mut self,
|
||||
input_batch: &RecordBatch,
|
||||
) -> Result<RecordBatch, arrow::error::ArrowError> {
|
||||
let num_rows = input_batch.num_rows();
|
||||
let mut token_builder = ListBuilder::new(Int64Builder::with_capacity(num_rows * 1024)); // Pre-allocate space
|
||||
let mut should_break = false;
|
||||
|
||||
let column = input_batch.column_by_name("text").unwrap();
|
||||
let string_array = column
|
||||
.as_any()
|
||||
.downcast_ref::<arrow::array::StringArray>()
|
||||
.unwrap();
|
||||
for i in 0..num_rows {
|
||||
if self.cur_samples_cnt >= self.num_samples {
|
||||
should_break = true;
|
||||
break;
|
||||
}
|
||||
if !Array::is_null(string_array, i) {
|
||||
let text = string_array.value(i);
|
||||
// Split paragraph into lines
|
||||
for line in text.split('\n') {
|
||||
if let Ok(encoding) = self.tokenizer.encode(line, true) {
|
||||
let tb_values = token_builder.values();
|
||||
for &id in encoding.get_ids() {
|
||||
tb_values.append_value(id as i64);
|
||||
}
|
||||
token_builder.append(true);
|
||||
self.cur_samples_cnt += 1;
|
||||
if self.cur_samples_cnt.is_multiple_of(5000) {
|
||||
println!("Processed {} rows", self.cur_samples_cnt);
|
||||
}
|
||||
if self.cur_samples_cnt >= self.num_samples {
|
||||
should_break = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Create array and shuffle it
|
||||
let input_ids_array = token_builder.finish();
|
||||
|
||||
// Create shuffled array by randomly sampling indices
|
||||
let mut rng = rand::rngs::StdRng::seed_from_u64(1337);
|
||||
let len = input_ids_array.len();
|
||||
let mut indices: Vec<u32> = (0..len as u32).collect();
|
||||
indices.shuffle(&mut rng);
|
||||
|
||||
// Take values in shuffled order
|
||||
let indices_array = UInt32Array::from(indices);
|
||||
let shuffled = arrow::compute::take(&input_ids_array, &indices_array, None)?;
|
||||
|
||||
let batch = RecordBatch::try_new(self.schema.clone(), vec![Arc::new(shuffled)]);
|
||||
if should_break {
|
||||
println!("Stop at {} rows", self.cur_samples_cnt);
|
||||
self.parquet_readers.clear();
|
||||
self.current_reader = None;
|
||||
}
|
||||
|
||||
batch
|
||||
}
|
||||
}
|
||||
|
||||
impl RecordBatchReader for WikiTextBatchReader {
|
||||
fn schema(&self) -> Arc<Schema> {
|
||||
self.schema.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl Iterator for WikiTextBatchReader {
|
||||
type Item = Result<RecordBatch, arrow::error::ArrowError>;
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
loop {
|
||||
// If we have a current reader, try to get next batch
|
||||
if let Some(reader) = &mut self.current_reader
|
||||
&& let Some(batch_result) = reader.next()
|
||||
{
|
||||
return Some(batch_result.and_then(|batch| self.process_batch(&batch)));
|
||||
}
|
||||
|
||||
// If no current reader or current reader is exhausted, try to get next reader
|
||||
if self.current_reader_idx < self.parquet_readers.len()
|
||||
&& let Some(builder) = self.parquet_readers[self.current_reader_idx].take()
|
||||
{
|
||||
match builder.build() {
|
||||
Ok(reader) => {
|
||||
self.current_reader = Some(Box::new(reader));
|
||||
self.current_reader_idx += 1;
|
||||
continue;
|
||||
}
|
||||
Err(e) => {
|
||||
return Some(Err(arrow::error::ArrowError::ExternalError(Box::new(e))));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// No more readers available
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
rt.block_on(async {
|
||||
// Load tokenizer
|
||||
let tokenizer = load_tokenizer("gpt2")?;
|
||||
|
||||
// Set up Hugging Face API
|
||||
// Download from https://huggingface.co/datasets/Salesforce/wikitext/tree/main/wikitext-103-raw-v1
|
||||
let api = Api::new()?;
|
||||
let repo = api.repo(Repo::with_revision(
|
||||
"Salesforce/wikitext".into(),
|
||||
RepoType::Dataset,
|
||||
"main".into(),
|
||||
));
|
||||
|
||||
// Define the parquet files we want to download
|
||||
let train_files = vec![
|
||||
"wikitext-103-raw-v1/train-00000-of-00002.parquet",
|
||||
"wikitext-103-raw-v1/train-00001-of-00002.parquet",
|
||||
];
|
||||
|
||||
let mut parquet_readers = Vec::new();
|
||||
for file in &train_files {
|
||||
println!("Downloading file: {}", file);
|
||||
let file_path = repo.get(file)?;
|
||||
let data = std::fs::read(file_path)?;
|
||||
|
||||
// Create a temporary file in the system temp directory and write the downloaded data to it
|
||||
let mut temp_file = NamedTempFile::new()?;
|
||||
temp_file.write_all(&data)?;
|
||||
|
||||
// Create the parquet reader builder with a larger batch size
|
||||
let builder = ParquetRecordBatchReaderBuilder::try_new(temp_file.into_file())?
|
||||
.with_batch_size(8192); // Increase batch size for better performance
|
||||
parquet_readers.push(builder);
|
||||
}
|
||||
|
||||
if parquet_readers.is_empty() {
|
||||
println!("No parquet files found to process.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Create batch reader
|
||||
let num_samples: u64 = 500_000;
|
||||
let batch_reader = WikiTextBatchReader::new(parquet_readers, tokenizer, Some(num_samples))?;
|
||||
|
||||
// Save as Lance dataset
|
||||
println!("Writing to Lance dataset...");
|
||||
let lance_dataset_path = "rust_wikitext_lance_dataset.lance";
|
||||
|
||||
let write_params = WriteParams::default();
|
||||
lance::Dataset::write(batch_reader, lance_dataset_path, Some(write_params)).await?;
|
||||
|
||||
// Verify the dataset
|
||||
let ds = lance::Dataset::open(lance_dataset_path).await?;
|
||||
let scanner = ds.scan();
|
||||
let mut stream = scanner.try_into_stream().await?;
|
||||
|
||||
let mut total_rows = 0;
|
||||
while let Some(batch_result) = stream.next().await {
|
||||
let batch = batch_result?;
|
||||
total_rows += batch.num_rows();
|
||||
}
|
||||
|
||||
println!(
|
||||
"Lance dataset created successfully with {} rows",
|
||||
total_rows
|
||||
);
|
||||
println!("Dataset location: {}", lance_dataset_path);
|
||||
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn load_tokenizer(model_name: &str) -> Result<Tokenizer, Box<dyn Error + Send + Sync>> {
|
||||
let api = Api::new()?;
|
||||
let repo = api.repo(Repo::with_revision(
|
||||
model_name.into(),
|
||||
RepoType::Model,
|
||||
"main".into(),
|
||||
));
|
||||
|
||||
let tokenizer_path = repo.get("tokenizer.json")?;
|
||||
let tokenizer = Tokenizer::from_file(tokenizer_path)?;
|
||||
|
||||
Ok(tokenizer)
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
#![allow(clippy::print_stdout)]
|
||||
|
||||
use arrow::array::UInt32Array;
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use arrow::record_batch::{RecordBatch, RecordBatchIterator};
|
||||
use futures::StreamExt;
|
||||
use lance::Dataset;
|
||||
use lance::dataset::{WriteMode, WriteParams};
|
||||
use lance::io::ObjectStore;
|
||||
use lance_core::utils::tempfile::TempStrDir;
|
||||
use std::sync::Arc;
|
||||
|
||||
// Writes sample dataset to the given path
|
||||
async fn write_dataset(data_path: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Define new schema
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("key", DataType::UInt32, false),
|
||||
Field::new("value", DataType::UInt32, false),
|
||||
]));
|
||||
|
||||
// Create new record batches
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6])),
|
||||
Arc::new(UInt32Array::from(vec![6, 7, 8, 9, 10, 11])),
|
||||
],
|
||||
)?;
|
||||
|
||||
let batches = RecordBatchIterator::new([Ok(batch)], schema.clone());
|
||||
|
||||
// Define write parameters (e.g. overwrite dataset)
|
||||
let write_params = WriteParams {
|
||||
mode: WriteMode::Overwrite,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
Dataset::write(batches, data_path, Some(write_params)).await?;
|
||||
Ok(())
|
||||
} // End write dataset
|
||||
|
||||
// Reads dataset from the given path and prints batch size, schema for all record batches. Also extracts and prints a slice from the first batch
|
||||
async fn read_dataset(data_path: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let dataset = Dataset::open(data_path).await?;
|
||||
let scanner = dataset.scan();
|
||||
|
||||
let mut batch_stream = scanner.try_into_stream().await?.map(|b| b.unwrap());
|
||||
|
||||
while let Some(batch) = batch_stream.next().await {
|
||||
println!("Batch size: {}, {}", batch.num_rows(), batch.num_columns()); // print size of batch
|
||||
println!("Schema: {:?}", batch.schema()); // print schema of recordbatch
|
||||
|
||||
println!("Batch: {:?}", batch); // print the entire recordbatch (schema and data)
|
||||
}
|
||||
Ok(())
|
||||
} // End read dataset
|
||||
|
||||
async fn clean_resources(data_path: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let (store, base) = ObjectStore::from_uri(data_path).await?;
|
||||
store.remove_dir_all(base).await?;
|
||||
println!("Cleaned up resources at: {}", data_path);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let tempdir = TempStrDir::default();
|
||||
let data_path = tempdir.as_str();
|
||||
|
||||
let result = async {
|
||||
write_dataset(data_path).await?;
|
||||
read_dataset(data_path).await?;
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
|
||||
if let Err(e) = clean_resources(data_path).await {
|
||||
eprintln!("Failed to clean resources: {:?}", e);
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 305 KiB |
@@ -0,0 +1,36 @@
|
||||
[package]
|
||||
name = "lance-arrow"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
readme = "README.md"
|
||||
description = "Arrow Extension for Lance"
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[dependencies]
|
||||
arrow-array = { workspace = true }
|
||||
arrow-buffer = { workspace = true }
|
||||
arrow-data = { workspace = true }
|
||||
arrow-ipc = { workspace = true }
|
||||
arrow-ord = { workspace = true }
|
||||
arrow-schema = { workspace = true }
|
||||
arrow-select = { workspace = true }
|
||||
bytemuck = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
half = { workspace = true }
|
||||
jsonb ={ workspace = true }
|
||||
num-traits = { workspace = true }
|
||||
rand.workspace = true
|
||||
|
||||
[target.'cfg(target_arch = "wasm32")'.dependencies]
|
||||
getrandom = { version = "0.2", features = ["js"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,8 @@
|
||||
# lance-arrow
|
||||
|
||||
`lance-arrow` is a internal sub-crate, containing [Apache-Arrow](https://github.com/apache/arrow-rs)
|
||||
extensions used by [Lance](https://github.com/lancedb/lance).
|
||||
|
||||
**Important Note**: This crate is **not intended for external usage**.
|
||||
|
||||
|
||||
@@ -0,0 +1,407 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! bfloat16 support for Apache Arrow.
|
||||
|
||||
use std::fmt::Formatter;
|
||||
use std::slice;
|
||||
|
||||
use arrow_array::{Array, FixedSizeBinaryArray, builder::BooleanBufferBuilder};
|
||||
use arrow_buffer::{Buffer, MutableBuffer};
|
||||
use arrow_data::ArrayData;
|
||||
use arrow_schema::{ArrowError, DataType, Field as ArrowField};
|
||||
use half::bf16;
|
||||
|
||||
use crate::{ARROW_EXT_NAME_KEY, FloatArray};
|
||||
|
||||
/// The name of the bfloat16 extension in Arrow metadata
|
||||
pub const BFLOAT16_EXT_NAME: &str = "lance.bfloat16";
|
||||
|
||||
/// Check whether the given field is a bfloat16 field
|
||||
///
|
||||
/// A field is a bfloat16 field if it has a data type of `FixedSizeBinary(2)` and the metadata
|
||||
/// contains the bfloat16 extension name.
|
||||
pub fn is_bfloat16_field(field: &ArrowField) -> bool {
|
||||
field.data_type() == &DataType::FixedSizeBinary(2)
|
||||
&& field
|
||||
.metadata()
|
||||
.get(ARROW_EXT_NAME_KEY)
|
||||
.map(|name| name == BFLOAT16_EXT_NAME)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// The bfloat16 data type
|
||||
///
|
||||
/// This implements the [`ArrowFloatType`](crate::floats::ArrowFloatType) trait for bfloat16 values.
|
||||
#[derive(Debug)]
|
||||
pub struct BFloat16Type {}
|
||||
|
||||
/// An array of bfloat16 values
|
||||
///
|
||||
/// Note that bfloat16 is not the same thing as fp16 which is supported natively by arrow-rs.
|
||||
#[derive(Clone)]
|
||||
pub struct BFloat16Array {
|
||||
inner: FixedSizeBinaryArray,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for BFloat16Array {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "BFloat16Array\n[\n")?;
|
||||
from_arrow::print_long_array(&self.inner, f, |array, i, f| {
|
||||
if array.is_null(i) {
|
||||
write!(f, "null")
|
||||
} else {
|
||||
let binary_values = array.value(i);
|
||||
let value =
|
||||
bf16::from_bits(u16::from_le_bytes([binary_values[0], binary_values[1]]));
|
||||
write!(f, "{:?}", value)
|
||||
}
|
||||
})?;
|
||||
write!(f, "]")
|
||||
}
|
||||
}
|
||||
|
||||
impl BFloat16Array {
|
||||
pub fn from_iter_values(iter: impl IntoIterator<Item = bf16>) -> Self {
|
||||
let values: Vec<bf16> = iter.into_iter().collect();
|
||||
values.into()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.inner.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.inner.is_empty()
|
||||
}
|
||||
|
||||
pub fn is_null(&self, i: usize) -> bool {
|
||||
self.inner.is_null(i)
|
||||
}
|
||||
|
||||
pub fn null_count(&self) -> usize {
|
||||
self.inner.null_count()
|
||||
}
|
||||
|
||||
pub fn iter(&self) -> BFloat16Iter<'_> {
|
||||
BFloat16Iter {
|
||||
array: self,
|
||||
index: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn value(&self, i: usize) -> bf16 {
|
||||
assert!(
|
||||
i < self.len(),
|
||||
"Trying to access an element at index {} from a BFloat16Array of length {}",
|
||||
i,
|
||||
self.len()
|
||||
);
|
||||
// Safety:
|
||||
// `i < self.len()
|
||||
unsafe { self.value_unchecked(i) }
|
||||
}
|
||||
|
||||
/// # Safety
|
||||
/// Caller must ensure that `i < self.len()`
|
||||
pub unsafe fn value_unchecked(&self, i: usize) -> bf16 {
|
||||
let binary_value = self.inner.value_unchecked(i);
|
||||
bf16::from_bits(u16::from_le_bytes([binary_value[0], binary_value[1]]))
|
||||
}
|
||||
|
||||
pub fn into_inner(self) -> FixedSizeBinaryArray {
|
||||
self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl FromIterator<Option<bf16>> for BFloat16Array {
|
||||
fn from_iter<I: IntoIterator<Item = Option<bf16>>>(iter: I) -> Self {
|
||||
let mut buffer = MutableBuffer::new(10);
|
||||
// No null buffer builder :(
|
||||
let mut nulls = BooleanBufferBuilder::new(10);
|
||||
let mut len = 0;
|
||||
|
||||
for maybe_value in iter {
|
||||
if let Some(value) = maybe_value {
|
||||
let bytes = value.to_le_bytes();
|
||||
buffer.extend(bytes);
|
||||
} else {
|
||||
buffer.extend([0u8, 0u8]);
|
||||
}
|
||||
nulls.append(maybe_value.is_some());
|
||||
len += 1;
|
||||
}
|
||||
|
||||
let null_buffer = nulls.finish();
|
||||
let num_valid = null_buffer.count_set_bits();
|
||||
let null_buffer = if num_valid == len {
|
||||
None
|
||||
} else {
|
||||
Some(null_buffer.into_inner())
|
||||
};
|
||||
|
||||
let array_data = ArrayData::builder(DataType::FixedSizeBinary(2))
|
||||
.len(len)
|
||||
.add_buffer(buffer.into())
|
||||
.null_bit_buffer(null_buffer);
|
||||
// SAFETY: the value buffer contains exactly `2 * len` bytes (two bytes
|
||||
// pushed per iteration of the loop above, including the zero-fill for
|
||||
// null slots), which matches the `FixedSizeBinary(2)` storage layout.
|
||||
// The null bit buffer, when present, has `len` bits appended above, so
|
||||
// its length covers the array's logical range.
|
||||
let array_data = unsafe { array_data.build_unchecked() };
|
||||
Self {
|
||||
inner: FixedSizeBinaryArray::from(array_data),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromIterator<bf16> for BFloat16Array {
|
||||
fn from_iter<I: IntoIterator<Item = bf16>>(iter: I) -> Self {
|
||||
Self::from_iter_values(iter)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Vec<bf16>> for BFloat16Array {
|
||||
fn from(data: Vec<bf16>) -> Self {
|
||||
let len = data.len();
|
||||
// Zero-copy: `bf16` is `#[repr(transparent)]` over `u16` and derives
|
||||
// `bytemuck::Pod`, so `cast_vec` reinterprets the allocation in place —
|
||||
// no per-element copy or heap alloc. The crate-root `compile_error!`
|
||||
// pins `target_endian = "little"`, so the resulting bytes match the
|
||||
// `FixedSizeBinary(2)` on-disk order Lance writes elsewhere.
|
||||
let raw: Vec<u16> = bytemuck::cast_vec(data);
|
||||
let array_data = ArrayData::builder(DataType::FixedSizeBinary(2))
|
||||
.len(len)
|
||||
.add_buffer(Buffer::from_vec(raw));
|
||||
// SAFETY: the value buffer contains exactly `2 * len` bytes — one
|
||||
// `u16` per element after the layout-compatible cast — matching the
|
||||
// `FixedSizeBinary(2)` storage layout. No null buffer is attached, so
|
||||
// every element is logically valid.
|
||||
let array_data = unsafe { array_data.build_unchecked() };
|
||||
Self {
|
||||
inner: FixedSizeBinaryArray::from(array_data),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<FixedSizeBinaryArray> for BFloat16Array {
|
||||
type Error = ArrowError;
|
||||
|
||||
fn try_from(value: FixedSizeBinaryArray) -> Result<Self, Self::Error> {
|
||||
if value.value_length() == 2 {
|
||||
Ok(Self { inner: value })
|
||||
} else {
|
||||
Err(ArrowError::InvalidArgumentError(
|
||||
"FixedSizeBinaryArray must have a value length of 2".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<Self> for BFloat16Array {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.inner.eq(&other.inner)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct BFloat16Iter<'a> {
|
||||
array: &'a BFloat16Array,
|
||||
index: usize,
|
||||
}
|
||||
|
||||
impl<'a> Iterator for BFloat16Iter<'a> {
|
||||
type Item = Option<bf16>;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
if self.index >= self.array.len() {
|
||||
return None;
|
||||
}
|
||||
let i = self.index;
|
||||
self.index += 1;
|
||||
if self.array.is_null(i) {
|
||||
Some(None)
|
||||
} else {
|
||||
Some(Some(self.array.value(i)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Methods that are lifted from arrow-rs temporarily until they are made public.
|
||||
mod from_arrow {
|
||||
use arrow_array::Array;
|
||||
|
||||
/// Helper function for printing potentially long arrays.
|
||||
pub(super) fn print_long_array<A, F>(
|
||||
array: &A,
|
||||
f: &mut std::fmt::Formatter,
|
||||
print_item: F,
|
||||
) -> std::fmt::Result
|
||||
where
|
||||
A: Array,
|
||||
F: Fn(&A, usize, &mut std::fmt::Formatter) -> std::fmt::Result,
|
||||
{
|
||||
let head = std::cmp::min(10, array.len());
|
||||
|
||||
for i in 0..head {
|
||||
if array.is_null(i) {
|
||||
writeln!(f, " null,")?;
|
||||
} else {
|
||||
write!(f, " ")?;
|
||||
print_item(array, i, f)?;
|
||||
writeln!(f, ",")?;
|
||||
}
|
||||
}
|
||||
if array.len() > 10 {
|
||||
if array.len() > 20 {
|
||||
writeln!(f, " ...{} elements...,", array.len() - 20)?;
|
||||
}
|
||||
|
||||
let tail = std::cmp::max(head, array.len() - 10);
|
||||
|
||||
for i in tail..array.len() {
|
||||
if array.is_null(i) {
|
||||
writeln!(f, " null,")?;
|
||||
} else {
|
||||
write!(f, " ")?;
|
||||
print_item(array, i, f)?;
|
||||
writeln!(f, ",")?;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl FloatArray<BFloat16Type> for FixedSizeBinaryArray {
|
||||
type FloatType = BFloat16Type;
|
||||
|
||||
/// Returns the underlying `bf16` values as a borrowed slice.
|
||||
///
|
||||
/// # Preconditions
|
||||
///
|
||||
/// - `value_length()` must be 2 (the `FixedSizeBinary(2)` storage shape
|
||||
/// used by [`BFloat16Array`]). Asserted at entry.
|
||||
/// - The value buffer must be at least 2-byte aligned. Lance's in-tree
|
||||
/// constructors always satisfy this: value buffers are built either via
|
||||
/// `MutableBuffer` (aligned to arrow-buffer's `ALIGNMENT` constant, ≥32
|
||||
/// bytes) or via `Buffer::from_vec::<u16>` (aligned to `align_of::<u16>()`
|
||||
/// == 2); both meet `bf16`'s 2-byte requirement. Externally-built
|
||||
/// `FixedSizeBinaryArray`s arriving via FFI, IPC, or
|
||||
/// `Buffer::from_custom_allocation` are not required by arrow-rs to be
|
||||
/// aligned beyond a single byte; passing one to this method violates the
|
||||
/// precondition. A `debug_assert` below catches such inputs in debug and
|
||||
/// test builds.
|
||||
///
|
||||
/// # Endianness
|
||||
///
|
||||
/// `lance-arrow` is gated on `target_endian = "little"` at the crate root,
|
||||
/// so this method always returns values in the same byte order Lance writes
|
||||
/// (see [`BFloat16Array::value`] and the [`FromIterator`] impls).
|
||||
fn as_slice(&self) -> &[bf16] {
|
||||
assert_eq!(
|
||||
self.value_length(),
|
||||
2,
|
||||
"BFloat16 arrays must use FixedSizeBinary(2) storage"
|
||||
);
|
||||
debug_assert_eq!(
|
||||
(self.value_data().as_ptr() as usize) % std::mem::align_of::<bf16>(),
|
||||
0,
|
||||
"BFloat16 value buffer must be at least 2-byte aligned"
|
||||
);
|
||||
// SAFETY:
|
||||
// - The assert above pins `value_size == 2`, so `value_data().len() / 2`
|
||||
// equals the array's logical element count.
|
||||
// `FixedSizeBinaryArray::From<ArrayData>` constructs its value buffer
|
||||
// as `buffers[0].slice_with_length(offset * 2, len * 2)` (arrow-array
|
||||
// `fixed_size_binary_array.rs`), so `value_data()` already returns
|
||||
// the offset-adjusted slice. Do not replace `value_data()` with an
|
||||
// accessor that returns the un-sliced backing buffer.
|
||||
// - `bf16` is `#[repr(transparent)]` over `u16` (size 2, alignment 2);
|
||||
// every `u16` bit pattern is a valid `bf16`, so any byte content
|
||||
// yields a defined value — never UB.
|
||||
// - Alignment is the caller's responsibility per the precondition
|
||||
// documented above. The `debug_assert_eq!` immediately preceding this
|
||||
// block catches violations in debug and test builds only — release
|
||||
// builds rely on callers honoring the precondition. arrow-rs
|
||||
// declares `FixedSizeBinary(n)`'s
|
||||
// `BufferSpec::FixedWidth { alignment: align_of::<u8>() == 1 }`
|
||||
// (arrow-data `data.rs`), so arrow-rs alone does not guarantee
|
||||
// 2-byte alignment. Lance's in-tree construction paths build value
|
||||
// buffers via `MutableBuffer` (arrow-buffer `ALIGNMENT` constant,
|
||||
// ≥32 bytes) or `Buffer::from_vec::<u16>` (2-byte aligned), both of
|
||||
// which satisfy `bf16`'s 2-byte requirement.
|
||||
// - The returned slice borrows from `self`; the underlying ref-counted,
|
||||
// immutable Arrow buffer cannot be mutated or freed for the slice's
|
||||
// lifetime.
|
||||
unsafe {
|
||||
slice::from_raw_parts(
|
||||
self.value_data().as_ptr() as *const bf16,
|
||||
self.value_data().len() / 2,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn from_values(values: Vec<bf16>) -> Self {
|
||||
BFloat16Array::from(values).into_inner()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_basics() {
|
||||
let values: Vec<f32> = vec![1.0, 2.0, 3.0];
|
||||
let values: Vec<bf16> = values.iter().map(|v| bf16::from_f32(*v)).collect();
|
||||
|
||||
let array = BFloat16Array::from_iter_values(values.clone());
|
||||
let array2 = BFloat16Array::from(values.clone());
|
||||
assert_eq!(array, array2);
|
||||
assert_eq!(array.len(), 3);
|
||||
|
||||
// Pin the raw little-endian bytes emitted by `From<Vec<bf16>>` (rewritten to
|
||||
// reinterpret the Vec via `bytemuck::cast_vec`), so a layout/byte-order
|
||||
// regression is caught directly rather than only through Debug formatting.
|
||||
// bf16 is the high 16 bits of the f32: 1.0->0x3F80, 2.0->0x4000, 3.0->0x4040.
|
||||
let inner = array2.clone().into_inner();
|
||||
let raw_bytes: Vec<u8> = (0..inner.len())
|
||||
.flat_map(|i| inner.value(i).to_vec())
|
||||
.collect();
|
||||
assert_eq!(raw_bytes, vec![0x80, 0x3F, 0x00, 0x40, 0x40, 0x40]);
|
||||
|
||||
let expected_fmt = "BFloat16Array\n[\n 1.0,\n 2.0,\n 3.0,\n]";
|
||||
assert_eq!(expected_fmt, format!("{:?}", array));
|
||||
|
||||
for (expected, value) in values.iter().zip(array.iter()) {
|
||||
assert_eq!(Some(*expected), value);
|
||||
}
|
||||
|
||||
for (expected, value) in values.as_slice().iter().zip(array2.iter()) {
|
||||
assert_eq!(Some(*expected), value);
|
||||
}
|
||||
|
||||
let arrow_array = array.into_inner();
|
||||
assert_eq!(arrow_array.as_slice(), values.as_slice());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nulls() {
|
||||
let values: Vec<Option<bf16>> =
|
||||
vec![Some(bf16::from_f32(1.0)), None, Some(bf16::from_f32(3.0))];
|
||||
let array = BFloat16Array::from_iter(values.clone());
|
||||
assert_eq!(array.len(), 3);
|
||||
assert_eq!(array.null_count(), 1);
|
||||
|
||||
let expected_fmt = "BFloat16Array\n[\n 1.0,\n null,\n 3.0,\n]";
|
||||
assert_eq!(expected_fmt, format!("{:?}", array));
|
||||
|
||||
for (expected, value) in values.iter().zip(array.iter()) {
|
||||
assert_eq!(*expected, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{Array, RecordBatch, make_array};
|
||||
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer};
|
||||
use arrow_data::{ArrayData, ArrayDataBuilder, transform::MutableArrayData};
|
||||
|
||||
pub fn deep_copy_buffer(buffer: &Buffer) -> Buffer {
|
||||
Buffer::from(buffer.as_slice())
|
||||
}
|
||||
|
||||
pub fn deep_copy_nulls(nulls: Option<&NullBuffer>) -> Option<NullBuffer> {
|
||||
let nulls = nulls?;
|
||||
let bit_buffer = deep_copy_buffer(nulls.inner().inner());
|
||||
// SAFETY: `null_count` is taken from the source `NullBuffer`, which already
|
||||
// upheld `NullBuffer::new_unchecked`'s invariant — the unset-bit count over
|
||||
// the logical bit slice `[bit_offset, bit_offset + bit_len)`. `NullBuffer::slice`
|
||||
// adjusts only `BooleanBuffer::bit_offset` / `bit_len` and never byte-advances
|
||||
// the inner `Buffer`, so `deep_copy_buffer` (which copies the source `Buffer`'s
|
||||
// `as_slice()` view from byte 0) reproduces the exact bit pattern at the same
|
||||
// bit offsets; the unset-bit count is therefore preserved. `BooleanBuffer::new`
|
||||
// panics (does not UB) if `bit_offset + bit_len > 8 * buffer.len()`, and the
|
||||
// copy has the same length, so that check still passes.
|
||||
Some(unsafe {
|
||||
NullBuffer::new_unchecked(
|
||||
BooleanBuffer::new(bit_buffer, nulls.offset(), nulls.len()),
|
||||
nulls.null_count(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn deep_copy_array_data(data: &ArrayData) -> ArrayData {
|
||||
let data_type = data.data_type().clone();
|
||||
let len = data.len();
|
||||
let nulls = deep_copy_nulls(data.nulls());
|
||||
let offset = data.offset();
|
||||
let buffers = data
|
||||
.buffers()
|
||||
.iter()
|
||||
.map(deep_copy_buffer)
|
||||
.collect::<Vec<_>>();
|
||||
let child_data = data
|
||||
.child_data()
|
||||
.iter()
|
||||
.map(deep_copy_array_data)
|
||||
.collect::<Vec<_>>();
|
||||
// SAFETY: `build_unchecked` inherits `ArrayData::new_unchecked`'s contract —
|
||||
// `(data_type, len, offset, nulls, buffers, child_data)` must form a valid
|
||||
// Arrow array. This call reproduces `data` structurally: `data_type`, `len`,
|
||||
// and `offset` are forwarded unchanged; each buffer is replaced by a byte-
|
||||
// identical copy of its offset-applied `as_slice()` view (the output buffer
|
||||
// is `MutableBuffer`-allocated, at least as aligned as the source); `nulls`
|
||||
// is deep-copied with the same bit offset/length and unset-bit count (see
|
||||
// `deep_copy_nulls`); `child_data` is recursively cloned with the same
|
||||
// guarantee. Every value-level invariant the source upheld — UTF-8 validity,
|
||||
// monotonic offsets, in-bounds dictionary indices, run-end monotonicity,
|
||||
// struct child-length matching — therefore transfers to the copy. If the
|
||||
// source `ArrayData` was itself constructed via `new_unchecked` with an
|
||||
// invalid payload, this function faithfully reproduces that invalidity.
|
||||
unsafe {
|
||||
ArrayDataBuilder::new(data_type)
|
||||
.len(len)
|
||||
.nulls(nulls)
|
||||
.offset(offset)
|
||||
.buffers(buffers)
|
||||
.child_data(child_data)
|
||||
.build_unchecked()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn deep_copy_array(array: &dyn Array) -> Arc<dyn Array> {
|
||||
let data = array.to_data();
|
||||
let data = deep_copy_array_data(&data);
|
||||
make_array(data)
|
||||
}
|
||||
|
||||
pub fn deep_copy_batch(batch: &RecordBatch) -> crate::Result<RecordBatch> {
|
||||
let arrays = batch
|
||||
.columns()
|
||||
.iter()
|
||||
.map(|array| deep_copy_array(array))
|
||||
.collect::<Vec<_>>();
|
||||
RecordBatch::try_new(batch.schema(), arrays)
|
||||
}
|
||||
|
||||
/// Deep copy array data, extracting only the sliced portion using MutableArrayData
|
||||
/// This is the most efficient and correct way to copy just the sliced data
|
||||
pub fn deep_copy_array_data_sliced(data: &ArrayData) -> ArrayData {
|
||||
// Use MutableArrayData to efficiently copy just the slice
|
||||
let mut mutable = MutableArrayData::new(vec![data], false, data.len());
|
||||
|
||||
// Copy from offset to offset+len (the visible slice)
|
||||
mutable.extend(0, data.offset(), data.offset() + data.len());
|
||||
|
||||
// Freeze into immutable ArrayData
|
||||
mutable.freeze()
|
||||
}
|
||||
|
||||
/// Deep copy an array, extracting only the sliced portion using MutableArrayData
|
||||
pub fn deep_copy_array_sliced(array: &dyn Array) -> Arc<dyn Array> {
|
||||
let data = array.to_data();
|
||||
let data = deep_copy_array_data_sliced(&data);
|
||||
make_array(data)
|
||||
}
|
||||
|
||||
/// Deep copy a RecordBatch, extracting only the sliced portion using MutableArrayData
|
||||
pub fn deep_copy_batch_sliced(batch: &RecordBatch) -> crate::Result<RecordBatch> {
|
||||
let arrays = batch
|
||||
.columns()
|
||||
.iter()
|
||||
.map(|array| deep_copy_array_sliced(array))
|
||||
.collect::<Vec<_>>();
|
||||
RecordBatch::try_new(batch.schema(), arrays)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{Array, Int32Array, RecordBatch, StringArray};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
|
||||
#[test]
|
||||
fn test_deep_copy_sliced_array_with_nulls() {
|
||||
let array = Arc::new(Int32Array::from(vec![
|
||||
Some(1),
|
||||
None,
|
||||
Some(3),
|
||||
None,
|
||||
Some(5),
|
||||
]));
|
||||
let sliced_array = array.slice(1, 3);
|
||||
let copied_array = super::deep_copy_array(&sliced_array);
|
||||
assert_eq!(sliced_array.len(), copied_array.len());
|
||||
assert_eq!(sliced_array.nulls(), copied_array.nulls());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deep_copy_array_data_sliced() {
|
||||
let array = Int32Array::from((0..1000).collect::<Vec<i32>>());
|
||||
let sliced = array.slice(100, 10);
|
||||
|
||||
let sliced_data = sliced.to_data();
|
||||
let copied_data = super::deep_copy_array_data_sliced(&sliced_data);
|
||||
|
||||
assert_eq!(copied_data.len(), 10);
|
||||
assert_eq!(copied_data.offset(), 0);
|
||||
|
||||
// Verify data correctness
|
||||
let copied_array = Int32Array::from(copied_data);
|
||||
for i in 0..10 {
|
||||
assert_eq!(copied_array.value(i), 100 + i as i32);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deep_copy_array_sliced() {
|
||||
let array = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
|
||||
let sliced = array.slice(1, 3);
|
||||
|
||||
let copied = super::deep_copy_array_sliced(&sliced);
|
||||
|
||||
assert_eq!(copied.len(), 3);
|
||||
let copied_int = copied.as_any().downcast_ref::<Int32Array>().unwrap();
|
||||
assert_eq!(copied_int.value(0), 2);
|
||||
assert_eq!(copied_int.value(1), 3);
|
||||
assert_eq!(copied_int.value(2), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deep_copy_batch_sliced() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int32, false),
|
||||
Field::new("name", DataType::Utf8, false),
|
||||
]));
|
||||
|
||||
let id_array = Arc::new(Int32Array::from((0..100).collect::<Vec<i32>>()));
|
||||
let name_array = Arc::new(StringArray::from(
|
||||
(0..100)
|
||||
.map(|i| format!("name_{}", i))
|
||||
.collect::<Vec<String>>(),
|
||||
));
|
||||
|
||||
let batch = RecordBatch::try_new(
|
||||
schema,
|
||||
vec![id_array as Arc<dyn Array>, name_array as Arc<dyn Array>],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let sliced = batch.slice(10, 5);
|
||||
let copied = super::deep_copy_batch_sliced(&sliced).unwrap();
|
||||
|
||||
assert_eq!(copied.num_rows(), 5);
|
||||
assert_eq!(copied.num_columns(), 2);
|
||||
|
||||
// Verify data correctness
|
||||
let id_col = copied
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<Int32Array>()
|
||||
.unwrap();
|
||||
let name_col = copied
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()
|
||||
.unwrap();
|
||||
|
||||
for i in 0..5 {
|
||||
assert_eq!(id_col.value(i), 10 + i as i32);
|
||||
assert_eq!(name_col.value(i), format!("name_{}", 10 + i));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deep_copy_array_sliced_with_nulls() {
|
||||
let array = Arc::new(Int32Array::from(vec![
|
||||
Some(1),
|
||||
None,
|
||||
Some(3),
|
||||
None,
|
||||
Some(5),
|
||||
]));
|
||||
let sliced = array.slice(1, 3); // [None, Some(3), None]
|
||||
|
||||
let copied = super::deep_copy_array_sliced(&sliced);
|
||||
|
||||
assert_eq!(copied.len(), 3);
|
||||
assert_eq!(copied.null_count(), 2); // Two nulls in the slice
|
||||
|
||||
let copied_int = copied.as_any().downcast_ref::<Int32Array>().unwrap();
|
||||
assert!(!copied_int.is_valid(0)); // None
|
||||
assert!(copied_int.is_valid(1)); // Some(3)
|
||||
assert!(!copied_int.is_valid(2)); // None
|
||||
assert_eq!(copied_int.value(1), 3);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Floats Array
|
||||
|
||||
use std::fmt::{Debug, Display};
|
||||
use std::iter::Sum;
|
||||
use std::sync::Arc;
|
||||
use std::{
|
||||
fmt::Formatter,
|
||||
ops::{AddAssign, DivAssign},
|
||||
};
|
||||
|
||||
use arrow_array::{
|
||||
Array, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array,
|
||||
types::{Float16Type, Float32Type, Float64Type},
|
||||
};
|
||||
use arrow_schema::{DataType, Field};
|
||||
use half::{bf16, f16};
|
||||
use num_traits::{AsPrimitive, Bounded, Float, FromPrimitive};
|
||||
|
||||
use super::bfloat16::{BFloat16Array, BFloat16Type};
|
||||
use crate::Result;
|
||||
use crate::bfloat16::is_bfloat16_field;
|
||||
|
||||
/// Float data type.
|
||||
///
|
||||
/// This helps differentiate between the different float types,
|
||||
/// because bf16 is not officially supported [DataType] in arrow-rs.
|
||||
#[derive(Debug)]
|
||||
pub enum FloatType {
|
||||
BFloat16,
|
||||
Float16,
|
||||
Float32,
|
||||
Float64,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for FloatType {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::BFloat16 => write!(f, "bfloat16"),
|
||||
Self::Float16 => write!(f, "float16"),
|
||||
Self::Float32 => write!(f, "float32"),
|
||||
Self::Float64 => write!(f, "float64"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to convert a [DataType] to a [FloatType]. To support bfloat16, always
|
||||
/// prefer using the `TryFrom<&Field>` implementation.
|
||||
impl TryFrom<&DataType> for FloatType {
|
||||
type Error = crate::ArrowError;
|
||||
|
||||
fn try_from(value: &DataType) -> Result<Self> {
|
||||
match *value {
|
||||
DataType::Float16 => Ok(Self::Float16),
|
||||
DataType::Float32 => Ok(Self::Float32),
|
||||
DataType::Float64 => Ok(Self::Float64),
|
||||
_ => Err(crate::ArrowError::InvalidArgumentError(format!(
|
||||
"{:?} is not a floating type",
|
||||
value
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&Field> for FloatType {
|
||||
type Error = crate::ArrowError;
|
||||
|
||||
fn try_from(field: &Field) -> Result<Self> {
|
||||
match field.data_type() {
|
||||
DataType::FixedSizeBinary(2) if is_bfloat16_field(field) => Ok(Self::BFloat16),
|
||||
_ => Self::try_from(field.data_type()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for float types used in Lance indexes
|
||||
///
|
||||
/// This mimics the utilities provided by [`arrow_array::ArrowPrimitiveType`]
|
||||
/// but applies to all float types (including bfloat16)
|
||||
pub trait ArrowFloatType: Debug {
|
||||
type Native: FromPrimitive
|
||||
+ FloatToArrayType<ArrowType = Self>
|
||||
+ AsPrimitive<f32>
|
||||
+ Debug
|
||||
+ Display;
|
||||
|
||||
const FLOAT_TYPE: FloatType;
|
||||
const MIN: Self::Native;
|
||||
const MAX: Self::Native;
|
||||
|
||||
/// Arrow Float Array Type.
|
||||
type ArrayType: FloatArray<Self>;
|
||||
|
||||
/// Returns empty array of this type.
|
||||
fn empty_array() -> Self::ArrayType {
|
||||
<Self::ArrayType as FloatArray<Self>>::from_values(Vec::new())
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait to be implemented by native types that have a corresponding [`ArrowFloatType`]
|
||||
/// implementation.
|
||||
///
|
||||
/// This helps define what operations are supported by native floats and also helps convert
|
||||
/// from a native type back to the corresponding Arrow float type.
|
||||
pub trait FloatToArrayType:
|
||||
Float
|
||||
+ Bounded
|
||||
+ Sum
|
||||
+ AddAssign<Self>
|
||||
+ AsPrimitive<f64>
|
||||
+ AsPrimitive<f32>
|
||||
+ DivAssign
|
||||
+ Send
|
||||
+ Sync
|
||||
+ Copy
|
||||
{
|
||||
/// The corresponding [`ArrowFloatType`] implementation for this native type
|
||||
type ArrowType: ArrowFloatType<Native = Self>;
|
||||
}
|
||||
|
||||
impl FloatToArrayType for bf16 {
|
||||
type ArrowType = BFloat16Type;
|
||||
}
|
||||
|
||||
impl FloatToArrayType for f16 {
|
||||
type ArrowType = Float16Type;
|
||||
}
|
||||
|
||||
impl FloatToArrayType for f32 {
|
||||
type ArrowType = Float32Type;
|
||||
}
|
||||
|
||||
impl FloatToArrayType for f64 {
|
||||
type ArrowType = Float64Type;
|
||||
}
|
||||
|
||||
impl ArrowFloatType for BFloat16Type {
|
||||
type Native = bf16;
|
||||
|
||||
const FLOAT_TYPE: FloatType = FloatType::BFloat16;
|
||||
const MIN: Self::Native = bf16::MIN;
|
||||
const MAX: Self::Native = bf16::MAX;
|
||||
|
||||
type ArrayType = FixedSizeBinaryArray;
|
||||
}
|
||||
|
||||
impl ArrowFloatType for Float16Type {
|
||||
type Native = f16;
|
||||
|
||||
const FLOAT_TYPE: FloatType = FloatType::Float16;
|
||||
const MIN: Self::Native = f16::MIN;
|
||||
const MAX: Self::Native = f16::MAX;
|
||||
|
||||
type ArrayType = Float16Array;
|
||||
}
|
||||
|
||||
impl ArrowFloatType for Float32Type {
|
||||
type Native = f32;
|
||||
|
||||
const FLOAT_TYPE: FloatType = FloatType::Float32;
|
||||
const MIN: Self::Native = f32::MIN;
|
||||
const MAX: Self::Native = f32::MAX;
|
||||
|
||||
type ArrayType = Float32Array;
|
||||
}
|
||||
|
||||
impl ArrowFloatType for Float64Type {
|
||||
type Native = f64;
|
||||
|
||||
const FLOAT_TYPE: FloatType = FloatType::Float64;
|
||||
const MIN: Self::Native = f64::MIN;
|
||||
const MAX: Self::Native = f64::MAX;
|
||||
|
||||
type ArrayType = Float64Array;
|
||||
}
|
||||
|
||||
/// [FloatArray] is a trait that is implemented by all float type arrays
|
||||
///
|
||||
/// This is similar to [`arrow_array::PrimitiveArray`] but applies to all float types (including bfloat16)
|
||||
/// and is implemented as a trait and not a struct
|
||||
pub trait FloatArray<T: ArrowFloatType + ?Sized>: Array + Clone + 'static {
|
||||
type FloatType: ArrowFloatType;
|
||||
|
||||
/// Returns a reference to the underlying data as a slice.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Implementations may panic if the array's storage shape does not match
|
||||
/// the expected element layout. In particular, the `bf16` impl panics if
|
||||
/// `value_length() != 2` (the `FixedSizeBinary(2)` shape required by
|
||||
/// `BFloat16Array`).
|
||||
///
|
||||
/// # Preconditions
|
||||
///
|
||||
/// Implementations may impose additional invariants on the underlying
|
||||
/// buffer. The `bf16` impl requires the value buffer to be at least
|
||||
/// 2-byte aligned — satisfied automatically by every in-tree Lance
|
||||
/// constructor, but external callers passing externally-built arrays
|
||||
/// (FFI, IPC, `Buffer::from_custom_allocation`) must ensure alignment.
|
||||
/// See the impl's docstring for details.
|
||||
fn as_slice(&self) -> &[T::Native];
|
||||
|
||||
/// Construct an array from a vector of values.
|
||||
fn from_values(values: Vec<T::Native>) -> Self;
|
||||
|
||||
/// Construct an array from an iterator of values.
|
||||
fn from_iter_values(values: impl IntoIterator<Item = T::Native>) -> Self
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
Self::from_values(values.into_iter().collect())
|
||||
}
|
||||
}
|
||||
|
||||
impl FloatArray<Float16Type> for Float16Array {
|
||||
type FloatType = Float16Type;
|
||||
|
||||
fn as_slice(&self) -> &[<Float16Type as ArrowFloatType>::Native] {
|
||||
self.values()
|
||||
}
|
||||
|
||||
fn from_values(values: Vec<<Float16Type as ArrowFloatType>::Native>) -> Self {
|
||||
Self::from(values)
|
||||
}
|
||||
}
|
||||
|
||||
impl FloatArray<Float32Type> for Float32Array {
|
||||
type FloatType = Float32Type;
|
||||
|
||||
fn as_slice(&self) -> &[<Float32Type as ArrowFloatType>::Native] {
|
||||
self.values()
|
||||
}
|
||||
|
||||
fn from_values(values: Vec<<Float32Type as ArrowFloatType>::Native>) -> Self {
|
||||
Self::from(values)
|
||||
}
|
||||
}
|
||||
|
||||
impl FloatArray<Float64Type> for Float64Array {
|
||||
type FloatType = Float64Type;
|
||||
|
||||
fn as_slice(&self) -> &[<Float64Type as ArrowFloatType>::Native] {
|
||||
self.values()
|
||||
}
|
||||
|
||||
fn from_values(values: Vec<<Float64Type as ArrowFloatType>::Native>) -> Self {
|
||||
Self::from(values)
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a float32 array to another float array
|
||||
///
|
||||
/// This is used during queries as query vectors are always provided as float32 arrays
|
||||
/// and need to be converted to the appropriate float type for the index.
|
||||
pub fn coerce_float_vector(input: &Float32Array, float_type: FloatType) -> Result<Arc<dyn Array>> {
|
||||
match float_type {
|
||||
FloatType::BFloat16 => Ok(Arc::new(
|
||||
BFloat16Array::from_iter_values(input.values().iter().map(|v| bf16::from_f32(*v)))
|
||||
.into_inner(),
|
||||
)),
|
||||
FloatType::Float16 => Ok(Arc::new(Float16Array::from_iter_values(
|
||||
input.values().iter().map(|v| f16::from_f32(*v)),
|
||||
))),
|
||||
FloatType::Float32 => Ok(Arc::new(input.clone())),
|
||||
FloatType::Float64 => Ok(Arc::new(Float64Array::from_iter_values(
|
||||
input.values().iter().map(|v| *v as f64),
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_coerce_float_vector_bfloat16() {
|
||||
let input = Float32Array::from(vec![1.0f32, 2.0, 3.0]);
|
||||
let array = coerce_float_vector(&input, FloatType::BFloat16).unwrap();
|
||||
|
||||
assert_eq!(array.data_type(), &DataType::FixedSizeBinary(2));
|
||||
|
||||
let fixed = array
|
||||
.as_any()
|
||||
.downcast_ref::<FixedSizeBinaryArray>()
|
||||
.unwrap();
|
||||
let expected: Vec<bf16> = input.values().iter().map(|v| bf16::from_f32(*v)).collect();
|
||||
assert_eq!(fixed.as_slice(), expected.as_slice());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,623 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Zero-copy Arrow IPC stream read/write utilities.
|
||||
//!
|
||||
//! Provides helpers for serializing and deserializing [`RecordBatch`]es as
|
||||
//! self-delimiting Arrow IPC streams using synchronous [`Read`]/[`Write`] I/O.
|
||||
//!
|
||||
//! These are designed for embedding IPC streams inside larger binary formats
|
||||
//! (e.g. a cache entry that contains multiple IPC sections). Each stream is
|
||||
//! self-delimiting (schema + batches + EOS marker) and can be read back
|
||||
//! independently.
|
||||
//!
|
||||
//! # Zero-copy reads
|
||||
//!
|
||||
//! [`read_ipc_stream`] and [`read_ipc_stream_single`] take `&Bytes` and use
|
||||
//! [`Bytes::slice`] to produce each message buffer. Because `Bytes::slice`
|
||||
//! increments a reference count rather than copying, the resulting
|
||||
//! [`Buffer`]s — and the array data decoded from them by [`FileDecoder`] —
|
||||
//! are all backed by the same allocation as the input.
|
||||
|
||||
use std::io::{Read, Write};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::RecordBatch;
|
||||
use arrow_buffer::Buffer;
|
||||
use arrow_ipc::convert::fb_to_schema;
|
||||
use arrow_ipc::reader::FileDecoder;
|
||||
use arrow_ipc::root_as_message;
|
||||
use arrow_ipc::writer::StreamWriter;
|
||||
use arrow_schema::ArrowError;
|
||||
use bytes::Bytes;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Length-prefixed byte utilities
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Write `data` prefixed by its length as a little-endian `u64`.
|
||||
///
|
||||
/// Paired with [`read_len_prefixed_bytes`].
|
||||
pub fn write_len_prefixed_bytes(writer: &mut dyn Write, data: &[u8]) -> Result<(), ArrowError> {
|
||||
writer
|
||||
.write_all(&(data.len() as u64).to_le_bytes())
|
||||
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
|
||||
writer
|
||||
.write_all(data)
|
||||
.map_err(|e| ArrowError::IoError(e.to_string(), e))
|
||||
}
|
||||
|
||||
/// Read a byte slice written by [`write_len_prefixed_bytes`].
|
||||
///
|
||||
/// Reads an 8-byte little-endian length then exactly that many bytes.
|
||||
pub fn read_len_prefixed_bytes(reader: &mut dyn Read) -> Result<Vec<u8>, ArrowError> {
|
||||
let mut len_buf = [0u8; 8];
|
||||
reader
|
||||
.read_exact(&mut len_buf)
|
||||
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
|
||||
let len = u64::from_le_bytes(len_buf) as usize;
|
||||
let mut buf = vec![0u8; len];
|
||||
reader
|
||||
.read_exact(&mut buf)
|
||||
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// IPC stream utilities
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// 4-byte continuation marker used by modern Arrow IPC streams.
|
||||
const IPC_CONTINUATION: [u8; 4] = [0xff; 4];
|
||||
|
||||
/// Write `batch` as a single-batch Arrow IPC stream to `writer`.
|
||||
pub fn write_ipc_stream(batch: &RecordBatch, writer: &mut dyn Write) -> Result<(), ArrowError> {
|
||||
let mut sw = StreamWriter::try_new(&mut *writer, batch.schema_ref())?;
|
||||
sw.write(batch)?;
|
||||
sw.finish()
|
||||
}
|
||||
|
||||
/// Write all batches from `iter` as a single Arrow IPC stream to `writer`.
|
||||
///
|
||||
/// `iter` must yield at least one batch; the schema is inferred from the first
|
||||
/// batch. Returns `ArrowError::InvalidArgumentError` if the iterator is empty.
|
||||
/// If you need to write an empty stream (schema only, no rows), construct a
|
||||
/// `StreamWriter` directly.
|
||||
pub fn write_ipc_stream_batches<I>(iter: I, writer: &mut dyn Write) -> Result<(), ArrowError>
|
||||
where
|
||||
I: IntoIterator<Item = RecordBatch>,
|
||||
{
|
||||
let mut iter = iter.into_iter();
|
||||
let first = iter
|
||||
.next()
|
||||
.ok_or_else(|| ArrowError::InvalidArgumentError("no batches to serialize".into()))?;
|
||||
let mut sw = StreamWriter::try_new(&mut *writer, first.schema_ref())?;
|
||||
sw.write(&first)?;
|
||||
for batch in iter {
|
||||
sw.write(&batch)?;
|
||||
}
|
||||
sw.finish()
|
||||
}
|
||||
|
||||
/// Read one complete Arrow IPC stream message from `data` as a zero-copy [`Buffer`].
|
||||
///
|
||||
/// Parses the first message starting at byte 0 of `data`. Returns `None` on
|
||||
/// EOS (size field == 0) or empty input. The returned [`Buffer`] is backed by
|
||||
/// `data`'s allocation — no bytes are copied.
|
||||
///
|
||||
/// The caller should advance its position by `buf.len()` after each call.
|
||||
fn read_one_ipc_message(data: &Bytes) -> Result<Option<Buffer>, ArrowError> {
|
||||
let bytes = data.as_ref();
|
||||
|
||||
if bytes.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
if bytes.len() < 4 {
|
||||
return Err(ArrowError::IoError(
|
||||
"IPC: truncated header".into(),
|
||||
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC header"),
|
||||
));
|
||||
}
|
||||
|
||||
let has_continuation = bytes[..4] == IPC_CONTINUATION;
|
||||
let (size_bytes, prefix_len): ([u8; 4], usize) = if has_continuation {
|
||||
if bytes.len() < 8 {
|
||||
return Err(ArrowError::IoError(
|
||||
"IPC: truncated header after continuation".into(),
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
"truncated after continuation",
|
||||
),
|
||||
));
|
||||
}
|
||||
(bytes[4..8].try_into().unwrap(), 8)
|
||||
} else {
|
||||
(bytes[..4].try_into().unwrap(), 4)
|
||||
};
|
||||
|
||||
let meta_size = u32::from_le_bytes(size_bytes) as usize;
|
||||
if meta_size == 0 {
|
||||
return Ok(None); // EOS
|
||||
}
|
||||
|
||||
let meta_end = prefix_len + meta_size;
|
||||
if bytes.len() < meta_end {
|
||||
return Err(ArrowError::IoError(
|
||||
"IPC: truncated metadata".into(),
|
||||
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC metadata"),
|
||||
));
|
||||
}
|
||||
|
||||
let msg = root_as_message(&bytes[prefix_len..meta_end])
|
||||
.map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?;
|
||||
let body_len = msg.bodyLength() as usize;
|
||||
|
||||
let total = meta_end + body_len;
|
||||
if bytes.len() < total {
|
||||
return Err(ArrowError::IoError(
|
||||
"IPC: truncated body".into(),
|
||||
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC body"),
|
||||
));
|
||||
}
|
||||
|
||||
// Zero-copy: Bytes::slice shares the backing allocation; Buffer::from
|
||||
// wraps it without copying.
|
||||
Ok(Some(Buffer::from(data.slice(0..total))))
|
||||
}
|
||||
|
||||
/// Read a length-prefixed byte slice at `offset` in `data`, advancing `offset`.
|
||||
///
|
||||
/// Reads an 8-byte little-endian length, then slices exactly that many bytes
|
||||
/// from `data`. The returned [`Bytes`] is zero-copy (shares `data`'s allocation).
|
||||
pub fn read_len_prefixed_bytes_at(data: &Bytes, offset: &mut usize) -> Result<Bytes, ArrowError> {
|
||||
let bytes = data.as_ref();
|
||||
let len_end = offset
|
||||
.checked_add(8)
|
||||
.filter(|&e| e <= bytes.len())
|
||||
.ok_or_else(|| {
|
||||
ArrowError::IoError(
|
||||
"length-prefixed bytes: truncated length field".into(),
|
||||
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated length"),
|
||||
)
|
||||
})?;
|
||||
let len = u64::from_le_bytes(bytes[*offset..len_end].try_into().unwrap()) as usize;
|
||||
*offset = len_end;
|
||||
let data_end = offset
|
||||
.checked_add(len)
|
||||
.filter(|&e| e <= bytes.len())
|
||||
.ok_or_else(|| {
|
||||
ArrowError::IoError(
|
||||
"length-prefixed bytes: truncated data".into(),
|
||||
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated data"),
|
||||
)
|
||||
})?;
|
||||
let result = data.slice(*offset..data_end);
|
||||
*offset = data_end;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Read all [`RecordBatch`]es from one Arrow IPC stream starting at `offset`,
|
||||
/// advancing `offset` past the stream (including the EOS marker).
|
||||
///
|
||||
/// Zero-copy: array buffers borrow from `data`'s allocation.
|
||||
pub fn read_ipc_stream_at(
|
||||
data: &Bytes,
|
||||
offset: &mut usize,
|
||||
) -> Result<Vec<RecordBatch>, ArrowError> {
|
||||
let batches = read_ipc_stream(&data.slice(*offset..))?;
|
||||
|
||||
// Recompute how many bytes were consumed by re-parsing message sizes.
|
||||
// We can't get this from read_ipc_stream directly, so we re-walk the
|
||||
// message headers (metadata only, no body re-read) to sum up lengths.
|
||||
let slice = &data.as_ref()[*offset..];
|
||||
let mut consumed = 0usize;
|
||||
loop {
|
||||
let rem = &slice[consumed..];
|
||||
if rem.is_empty() {
|
||||
break;
|
||||
}
|
||||
let has_cont = rem.len() >= 4 && rem[..4] == IPC_CONTINUATION;
|
||||
let (size_bytes, prefix_len): ([u8; 4], usize) = if has_cont {
|
||||
if rem.len() < 8 {
|
||||
break;
|
||||
}
|
||||
(rem[4..8].try_into().unwrap(), 8)
|
||||
} else {
|
||||
if rem.len() < 4 {
|
||||
break;
|
||||
}
|
||||
(rem[..4].try_into().unwrap(), 4)
|
||||
};
|
||||
let meta_size = u32::from_le_bytes(size_bytes) as usize;
|
||||
if meta_size == 0 {
|
||||
// EOS — consume it and stop.
|
||||
consumed += prefix_len;
|
||||
break;
|
||||
}
|
||||
let meta_end = prefix_len + meta_size;
|
||||
if rem.len() < meta_end {
|
||||
break;
|
||||
}
|
||||
let msg = root_as_message(&rem[prefix_len..meta_end])
|
||||
.map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?;
|
||||
let body_len = msg.bodyLength() as usize;
|
||||
consumed += meta_end + body_len;
|
||||
}
|
||||
*offset += consumed;
|
||||
|
||||
Ok(batches)
|
||||
}
|
||||
|
||||
/// Read exactly one [`RecordBatch`] from one Arrow IPC stream starting at `offset`,
|
||||
/// advancing `offset` past the stream (including the EOS marker).
|
||||
///
|
||||
/// Zero-copy: array buffers borrow from `data`'s allocation.
|
||||
pub fn read_ipc_stream_single_at(
|
||||
data: &Bytes,
|
||||
offset: &mut usize,
|
||||
) -> Result<RecordBatch, ArrowError> {
|
||||
let mut batches = read_ipc_stream_at(data, offset)?;
|
||||
match batches.len() {
|
||||
1 => Ok(batches.remove(0)),
|
||||
n => Err(ArrowError::ParseError(format!(
|
||||
"expected exactly 1 IPC record batch, got {n}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract the prefix length and metadata size from a raw IPC message buffer.
|
||||
///
|
||||
/// Modern IPC streams have an 8-byte prefix `[continuation: 4][size: 4]`.
|
||||
/// Legacy streams have a 4-byte prefix `[size: 4]`. Returns `(prefix_len, meta_size)`.
|
||||
fn parse_ipc_message_prefix(buf: &Buffer) -> Result<(usize, usize), ArrowError> {
|
||||
let has_continuation = buf.len() >= 4 && buf[..4] == IPC_CONTINUATION;
|
||||
if has_continuation {
|
||||
if buf.len() < 8 {
|
||||
return Err(ArrowError::ParseError(
|
||||
"IPC message buffer too short".into(),
|
||||
));
|
||||
}
|
||||
let meta_size = u32::from_le_bytes(buf[4..8].try_into().unwrap()) as usize;
|
||||
Ok((8, meta_size))
|
||||
} else {
|
||||
if buf.len() < 4 {
|
||||
return Err(ArrowError::ParseError(
|
||||
"IPC message buffer too short".into(),
|
||||
));
|
||||
}
|
||||
let meta_size = u32::from_le_bytes(buf[..4].try_into().unwrap()) as usize;
|
||||
Ok((4, meta_size))
|
||||
}
|
||||
}
|
||||
|
||||
/// Read all [`RecordBatch`]es from one Arrow IPC stream.
|
||||
///
|
||||
/// Zero-copy: each batch's array data buffers are borrowed from the input
|
||||
/// message buffer(s) and not copied during decoding.
|
||||
///
|
||||
/// Uses [`FileDecoder`] directly (rather than `StreamDecoder`) to avoid a
|
||||
/// known edge case where `StreamDecoder` does not produce a batch for messages
|
||||
/// with a zero-length body when the message exactly fills the decode buffer.
|
||||
pub fn read_ipc_stream(data: &Bytes) -> Result<Vec<RecordBatch>, ArrowError> {
|
||||
let mut offset = 0usize;
|
||||
|
||||
let schema_buf = read_one_ipc_message(&data.slice(offset..))?.ok_or_else(|| {
|
||||
ArrowError::ParseError("IPC stream: expected schema message, got EOS".into())
|
||||
})?;
|
||||
offset += schema_buf.len();
|
||||
|
||||
let (prefix_len, meta_size) = parse_ipc_message_prefix(&schema_buf)?;
|
||||
let schema_msg = root_as_message(&schema_buf[prefix_len..prefix_len + meta_size])
|
||||
.map_err(|e| ArrowError::ParseError(format!("IPC schema parse error: {e}")))?;
|
||||
let schema = Arc::new(fb_to_schema(schema_msg.header_as_schema().ok_or_else(
|
||||
|| ArrowError::ParseError("IPC stream: first message is not a schema".into()),
|
||||
)?));
|
||||
let mut decoder = FileDecoder::new(schema, schema_msg.version());
|
||||
|
||||
let mut batches = Vec::new();
|
||||
|
||||
loop {
|
||||
let Some(buf) = read_one_ipc_message(&data.slice(offset..))? else {
|
||||
break;
|
||||
};
|
||||
offset += buf.len();
|
||||
|
||||
let (prefix_len, meta_size) = parse_ipc_message_prefix(&buf)?;
|
||||
let msg = root_as_message(&buf[prefix_len..prefix_len + meta_size])
|
||||
.map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?;
|
||||
let body_len = msg.bodyLength() as usize;
|
||||
|
||||
// Block offset = 0 since the buffer starts at the message boundary.
|
||||
// metaDataLength = prefix_len + meta_size (prefix + flatbuf + padding).
|
||||
let block = arrow_ipc::Block::new(0, (prefix_len + meta_size) as i32, body_len as i64);
|
||||
|
||||
match msg.header_type() {
|
||||
arrow_ipc::MessageHeader::RecordBatch => {
|
||||
if let Some(batch) = decoder.read_record_batch(&block, &buf)? {
|
||||
batches.push(batch);
|
||||
}
|
||||
}
|
||||
arrow_ipc::MessageHeader::DictionaryBatch => {
|
||||
decoder.read_dictionary(&block, &buf)?;
|
||||
}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
|
||||
Ok(batches)
|
||||
}
|
||||
|
||||
/// Read exactly one [`RecordBatch`] from one Arrow IPC stream.
|
||||
pub fn read_ipc_stream_single(data: &Bytes) -> Result<RecordBatch, ArrowError> {
|
||||
let mut batches = read_ipc_stream(data)?;
|
||||
match batches.len() {
|
||||
1 => Ok(batches.remove(0)),
|
||||
n => Err(ArrowError::ParseError(format!(
|
||||
"expected exactly 1 IPC record batch, got {n}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Aligned IPC sections
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Byte alignment that each IPC section's stream start is padded to.
|
||||
///
|
||||
/// When several IPC streams are concatenated into one larger blob (e.g. a
|
||||
/// cache entry), a section that starts at an arbitrary offset would leave its
|
||||
/// array data misaligned. [`FileDecoder`] with `require_alignment = false`
|
||||
/// then silently copies each buffer into a freshly aligned allocation on
|
||||
/// every read, defeating zero-copy. Padding each section start to a 64-byte
|
||||
/// boundary keeps the decoded buffers borrowed directly from the input.
|
||||
pub const IPC_SECTION_ALIGNMENT: usize = 64;
|
||||
|
||||
/// Number of zero-padding bytes needed to advance `pos` to the next
|
||||
/// [`IPC_SECTION_ALIGNMENT`] boundary.
|
||||
fn section_padding(pos: usize) -> usize {
|
||||
(IPC_SECTION_ALIGNMENT - (pos % IPC_SECTION_ALIGNMENT)) % IPC_SECTION_ALIGNMENT
|
||||
}
|
||||
|
||||
/// A [`Write`] adapter that counts the bytes written through it.
|
||||
struct CountingWriter<'a> {
|
||||
inner: &'a mut dyn Write,
|
||||
count: usize,
|
||||
}
|
||||
|
||||
impl Write for CountingWriter<'_> {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
let n = self.inner.write(buf)?;
|
||||
self.count += n;
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
self.inner.flush()
|
||||
}
|
||||
}
|
||||
|
||||
/// Write zero padding so the next byte lands on an [`IPC_SECTION_ALIGNMENT`]
|
||||
/// boundary, advancing `pos` past it.
|
||||
fn write_section_padding(writer: &mut dyn Write, pos: &mut usize) -> Result<(), ArrowError> {
|
||||
let pad = section_padding(*pos);
|
||||
if pad > 0 {
|
||||
const ZEROS: [u8; IPC_SECTION_ALIGNMENT] = [0u8; IPC_SECTION_ALIGNMENT];
|
||||
writer
|
||||
.write_all(&ZEROS[..pad])
|
||||
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
|
||||
*pos += pad;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Write `batch` as a 64-byte-aligned single-batch Arrow IPC section.
|
||||
///
|
||||
/// `pos` is the absolute byte offset of `writer` within the enclosing blob.
|
||||
/// Zero padding is written first so the IPC stream begins on an
|
||||
/// [`IPC_SECTION_ALIGNMENT`] boundary, then the stream itself. `pos` is
|
||||
/// advanced past both the padding and the stream so the caller can write
|
||||
/// further aligned sections.
|
||||
///
|
||||
/// Paired with [`read_ipc_section_at`]. For the decoded buffers to be borrowed
|
||||
/// zero-copy, the blob must ultimately be read back from a buffer whose base
|
||||
/// address is at least 64-byte aligned.
|
||||
pub fn write_ipc_section(
|
||||
writer: &mut dyn Write,
|
||||
pos: &mut usize,
|
||||
batch: &RecordBatch,
|
||||
) -> Result<(), ArrowError> {
|
||||
write_section_padding(writer, pos)?;
|
||||
|
||||
let mut counting = CountingWriter {
|
||||
inner: writer,
|
||||
count: 0,
|
||||
};
|
||||
write_ipc_stream(batch, &mut counting)?;
|
||||
*pos += counting.count;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read a single [`RecordBatch`] from an aligned IPC section at `offset`.
|
||||
///
|
||||
/// Skips the alignment padding written by [`write_ipc_section`], then reads
|
||||
/// the stream, advancing `offset` past the section (padding + stream + EOS).
|
||||
///
|
||||
/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base
|
||||
/// address is at least 64-byte aligned (see [`write_ipc_section`]).
|
||||
pub fn read_ipc_section_at(data: &Bytes, offset: &mut usize) -> Result<RecordBatch, ArrowError> {
|
||||
*offset += section_padding(*offset);
|
||||
read_ipc_stream_single_at(data, offset)
|
||||
}
|
||||
|
||||
/// Write `batches` as a single 64-byte-aligned multi-batch Arrow IPC section.
|
||||
///
|
||||
/// Like [`write_ipc_section`] but emits every batch from `iter` into one IPC
|
||||
/// stream (schema + N batches + EOS). `iter` must yield at least one batch.
|
||||
/// Paired with [`read_ipc_section_batches_at`].
|
||||
pub fn write_ipc_section_batches<I>(
|
||||
writer: &mut dyn Write,
|
||||
pos: &mut usize,
|
||||
iter: I,
|
||||
) -> Result<(), ArrowError>
|
||||
where
|
||||
I: IntoIterator<Item = RecordBatch>,
|
||||
{
|
||||
write_section_padding(writer, pos)?;
|
||||
|
||||
let mut counting = CountingWriter {
|
||||
inner: writer,
|
||||
count: 0,
|
||||
};
|
||||
write_ipc_stream_batches(iter, &mut counting)?;
|
||||
*pos += counting.count;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read all [`RecordBatch`]es from an aligned multi-batch IPC section at
|
||||
/// `offset`, advancing `offset` past the section (padding + stream + EOS).
|
||||
///
|
||||
/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base
|
||||
/// address is at least 64-byte aligned (see [`write_ipc_section_batches`]).
|
||||
pub fn read_ipc_section_batches_at(
|
||||
data: &Bytes,
|
||||
offset: &mut usize,
|
||||
) -> Result<Vec<RecordBatch>, ArrowError> {
|
||||
*offset += section_padding(*offset);
|
||||
read_ipc_stream_at(data, offset)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use arrow_array::{ArrayRef, record_batch};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_ipc_roundtrip() {
|
||||
let batch1 = record_batch!(
|
||||
("int", Int32, [1, 2, 3]),
|
||||
("str", Utf8, ["foo", "bar", "baz"])
|
||||
)
|
||||
.unwrap();
|
||||
let batch2 = record_batch!(("int", Int32, [4, 5]), ("str", Utf8, ["qux", "quux"])).unwrap();
|
||||
let batches = vec![batch1.clone(), batch2.clone()];
|
||||
|
||||
let mut buf = Vec::new();
|
||||
write_ipc_stream_batches(batches, &mut buf).unwrap();
|
||||
|
||||
let data = Bytes::from(buf);
|
||||
|
||||
let batches = read_ipc_stream(&data).unwrap();
|
||||
assert_eq!(batches.len(), 2);
|
||||
assert_eq!(batches[0], batch1);
|
||||
assert_eq!(batches[1], batch2);
|
||||
|
||||
let data_base = data.as_ptr() as usize;
|
||||
let data_end = data_base + data.len();
|
||||
let assert_col_zero_copy = |array: &ArrayRef| {
|
||||
for buffer in array.to_data().buffers() {
|
||||
let ptr = buffer.as_ptr() as usize;
|
||||
assert!(
|
||||
ptr >= data_base && ptr < data_end,
|
||||
"buffer at {ptr:#x} is not backed by the input Bytes allocation \
|
||||
[{data_base:#x}..{data_end:#x})"
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
for batch in &batches {
|
||||
assert_eq!(batch.schema(), batch1.schema());
|
||||
assert_col_zero_copy(batch.column(0));
|
||||
assert_col_zero_copy(batch.column(1));
|
||||
}
|
||||
}
|
||||
|
||||
/// Allocate a [`Bytes`] whose base address is 64-byte aligned, modelling a
|
||||
/// backend that reads cache entries into an aligned buffer. A plain
|
||||
/// `Bytes::from(vec)` only guarantees the allocator's alignment for `u8`.
|
||||
fn aligned_bytes(payload: &[u8]) -> Bytes {
|
||||
let mut v = vec![0u8; payload.len() + IPC_SECTION_ALIGNMENT];
|
||||
let pad = section_padding(v.as_ptr() as usize);
|
||||
v[pad..pad + payload.len()].copy_from_slice(payload);
|
||||
Bytes::from(v).slice(pad..pad + payload.len())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_aligned_ipc_sections_are_zero_copy() {
|
||||
// A LargeBinary column exercises the i64-offset buffer whose 8-byte
|
||||
// alignment requirement triggers a realigning memcpy when misaligned.
|
||||
let blocks = arrow_array::LargeBinaryArray::from_vec(vec![&b"hello"[..], b"world"]);
|
||||
let section_a = RecordBatch::try_from_iter([("a", Arc::new(blocks) as ArrayRef)]).unwrap();
|
||||
let section_b = record_batch!(("b", Int64, [10i64, 20, 30, 40, 50])).unwrap();
|
||||
|
||||
let mut buf = Vec::new();
|
||||
// Arbitrary, deliberately non-64-aligned preamble so the first section
|
||||
// must be padded rather than landing at offset 0 by luck.
|
||||
buf.extend_from_slice(&[0xABu8; 7]);
|
||||
let mut pos = buf.len();
|
||||
// The first section's stream begins after padding the 7-byte preamble
|
||||
// up to the next 64-byte boundary.
|
||||
assert_eq!(7 + section_padding(7), IPC_SECTION_ALIGNMENT);
|
||||
write_ipc_section(&mut buf, &mut pos, §ion_a).unwrap();
|
||||
write_ipc_section(&mut buf, &mut pos, §ion_b).unwrap();
|
||||
|
||||
let data = aligned_bytes(&buf);
|
||||
assert_eq!(
|
||||
section_padding(data.as_ptr() as usize),
|
||||
0,
|
||||
"base not aligned"
|
||||
);
|
||||
|
||||
let mut offset = 7;
|
||||
let read_a = read_ipc_section_at(&data, &mut offset).unwrap();
|
||||
let read_b = read_ipc_section_at(&data, &mut offset).unwrap();
|
||||
assert_eq!(read_a, section_a);
|
||||
assert_eq!(read_b, section_b);
|
||||
|
||||
let data_base = data.as_ptr() as usize;
|
||||
let data_end = data_base + data.len();
|
||||
for batch in [&read_a, &read_b] {
|
||||
for buffer in batch.column(0).to_data().buffers() {
|
||||
let ptr = buffer.as_ptr() as usize;
|
||||
assert!(
|
||||
ptr >= data_base && ptr < data_end,
|
||||
"section buffer at {ptr:#x} was realigned out of the input \
|
||||
[{data_base:#x}..{data_end:#x}) — misaligned section",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_aligned_multi_batch_section_roundtrip_zero_copy() {
|
||||
// A multi-batch section (e.g. IVF SQ storage chunks) must round-trip
|
||||
// every batch and decode the first batch's buffers zero-copy.
|
||||
let b1 = record_batch!(("v", Int64, [1i64, 2, 3])).unwrap();
|
||||
let b2 = record_batch!(("v", Int64, [4i64, 5])).unwrap();
|
||||
let b3 = record_batch!(("v", Int64, [6i64])).unwrap();
|
||||
|
||||
let mut buf = vec![0xCDu8; 5];
|
||||
let mut pos = buf.len();
|
||||
write_ipc_section_batches(&mut buf, &mut pos, [b1.clone(), b2.clone(), b3.clone()])
|
||||
.unwrap();
|
||||
|
||||
let data = aligned_bytes(&buf);
|
||||
let mut offset = 5;
|
||||
let read = read_ipc_section_batches_at(&data, &mut offset).unwrap();
|
||||
assert_eq!(read, vec![b1, b2, b3]);
|
||||
assert_eq!(offset, buf.len(), "offset should land at section end");
|
||||
|
||||
let data_base = data.as_ptr() as usize;
|
||||
let data_end = data_base + data.len();
|
||||
for buffer in read[0].column(0).to_data().buffers() {
|
||||
let ptr = buffer.as_ptr() as usize;
|
||||
assert!(
|
||||
ptr >= data_base && ptr < data_end,
|
||||
"first batch buffer at {ptr:#x} was realigned out of the input",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,179 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{Array, BooleanArray, GenericListArray, OffsetSizeTrait};
|
||||
use arrow_buffer::{BooleanBufferBuilder, OffsetBuffer, ScalarBuffer};
|
||||
use arrow_schema::Field;
|
||||
|
||||
pub trait ListArrayExt {
|
||||
/// Filters out masked null items from the list array
|
||||
///
|
||||
/// It is legal for a list array to have a null entry with a non-zero length. The
|
||||
/// values inside the entry are "garbage" and should be ignored. This function
|
||||
/// filters the values array to remove the garbage values.
|
||||
///
|
||||
/// The output list will always have zero-length nulls.
|
||||
fn filter_garbage_nulls(&self) -> Self;
|
||||
/// Returns a copy of the list's values array that has been sliced to size
|
||||
///
|
||||
/// It is legal for a list array's offsets to not start with zero. It's also legal
|
||||
/// for a list array's offsets to not extend to the entire values array. This function
|
||||
/// behaves similarly to `values()` except it slices the array so that it starts at
|
||||
/// the first list offset and ends at the last list offset.
|
||||
fn trimmed_values(&self) -> Arc<dyn Array>;
|
||||
/// The offset type of the underlying list array.
|
||||
type Offset: OffsetSizeTrait;
|
||||
/// Returns offsets shifted so the first offset is zero, matching
|
||||
/// [`Self::trimmed_values`].
|
||||
///
|
||||
/// Sliced list arrays (e.g. a filtered batch) keep offsets that reference the
|
||||
/// original values buffer, so combining them with trimmed values produces
|
||||
/// offsets that exceed the values length. Use this together with
|
||||
/// `trimmed_values` when constructing a new list array.
|
||||
fn trimmed_offsets(&self) -> OffsetBuffer<Self::Offset>;
|
||||
}
|
||||
|
||||
impl<OffsetSize: OffsetSizeTrait> ListArrayExt for GenericListArray<OffsetSize> {
|
||||
fn filter_garbage_nulls(&self) -> Self {
|
||||
if self.is_empty() {
|
||||
return self.clone();
|
||||
}
|
||||
let Some(validity) = self.nulls().cloned() else {
|
||||
return self.clone();
|
||||
};
|
||||
|
||||
let mut should_keep = BooleanBufferBuilder::new(self.values().len());
|
||||
|
||||
// Handle case where offsets do not start at 0
|
||||
let preamble_len = self.offsets().first().unwrap().to_usize().unwrap();
|
||||
should_keep.append_n(preamble_len, false);
|
||||
|
||||
let mut new_offsets: Vec<OffsetSize> = Vec::with_capacity(self.len() + 1);
|
||||
new_offsets.push(OffsetSize::zero());
|
||||
let mut cur_len = OffsetSize::zero();
|
||||
for (offset, is_valid) in self.offsets().windows(2).zip(validity.iter()) {
|
||||
let len = offset[1] - offset[0];
|
||||
if is_valid {
|
||||
cur_len += len;
|
||||
should_keep.append_n(len.to_usize().unwrap(), true);
|
||||
new_offsets.push(cur_len);
|
||||
} else {
|
||||
should_keep.append_n(len.to_usize().unwrap(), false);
|
||||
new_offsets.push(cur_len);
|
||||
}
|
||||
}
|
||||
|
||||
// Offsets may not reference entire values buffer
|
||||
let trailer = self.values().len() - should_keep.len();
|
||||
should_keep.append_n(trailer, false);
|
||||
|
||||
let should_keep = should_keep.finish();
|
||||
let should_keep = BooleanArray::new(should_keep, None);
|
||||
let new_values = arrow_select::filter::filter(self.values(), &should_keep).unwrap();
|
||||
let new_offsets = ScalarBuffer::from(new_offsets);
|
||||
let new_offsets = OffsetBuffer::new(new_offsets);
|
||||
|
||||
Self::new(
|
||||
Arc::new(Field::new(
|
||||
"item",
|
||||
self.value_type(),
|
||||
self.values().is_nullable(),
|
||||
)),
|
||||
new_offsets,
|
||||
new_values,
|
||||
Some(validity),
|
||||
)
|
||||
}
|
||||
|
||||
fn trimmed_values(&self) -> Arc<dyn Array> {
|
||||
let first_value = self
|
||||
.offsets()
|
||||
.first()
|
||||
.map(|v| v.to_usize().unwrap())
|
||||
.unwrap_or(0);
|
||||
let last_value = self
|
||||
.offsets()
|
||||
.last()
|
||||
.map(|v| v.to_usize().unwrap())
|
||||
.unwrap_or(0);
|
||||
self.values().slice(first_value, last_value - first_value)
|
||||
}
|
||||
|
||||
type Offset = OffsetSize;
|
||||
|
||||
fn trimmed_offsets(&self) -> OffsetBuffer<OffsetSize> {
|
||||
let offsets = self.offsets();
|
||||
let Some(&first) = offsets.first() else {
|
||||
return offsets.clone();
|
||||
};
|
||||
if first == OffsetSize::zero() {
|
||||
return offsets.clone();
|
||||
}
|
||||
let shifted: Vec<OffsetSize> = offsets.iter().map(|&o| o - first).collect();
|
||||
OffsetBuffer::new(ScalarBuffer::from(shifted))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{ListArray, UInt64Array};
|
||||
use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
|
||||
use arrow_schema::{DataType, Field};
|
||||
|
||||
use super::ListArrayExt;
|
||||
|
||||
#[test]
|
||||
fn test_filter_garbage_nulls() {
|
||||
let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
|
||||
let offsets = ScalarBuffer::<i32>::from(vec![2, 5, 8, 9]);
|
||||
let offsets = OffsetBuffer::new(offsets);
|
||||
let list_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
|
||||
let list_arr = ListArray::new(
|
||||
Arc::new(Field::new("item", DataType::UInt64, true)),
|
||||
offsets,
|
||||
Arc::new(items),
|
||||
Some(list_validity.clone()),
|
||||
);
|
||||
|
||||
let filtered = list_arr.filter_garbage_nulls();
|
||||
|
||||
let expected_items = UInt64Array::from(vec![2, 3, 4, 8]);
|
||||
let offsets = ScalarBuffer::<i32>::from(vec![0, 3, 3, 4]);
|
||||
let expected = ListArray::new(
|
||||
Arc::new(Field::new("item", DataType::UInt64, false)),
|
||||
OffsetBuffer::new(offsets),
|
||||
Arc::new(expected_items),
|
||||
Some(list_validity),
|
||||
);
|
||||
|
||||
assert_eq!(filtered, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trim_values() {
|
||||
let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
|
||||
let offsets = ScalarBuffer::<i32>::from(vec![2, 5, 6, 8, 9]);
|
||||
let offsets = OffsetBuffer::new(offsets);
|
||||
let list_arr = ListArray::new(
|
||||
Arc::new(Field::new("item", DataType::UInt64, true)),
|
||||
offsets,
|
||||
Arc::new(items),
|
||||
None,
|
||||
);
|
||||
let list_arr = list_arr.slice(1, 2);
|
||||
|
||||
let trimmed = list_arr.trimmed_values();
|
||||
|
||||
let expected_items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
|
||||
let expected_items = expected_items.slice(5, 3);
|
||||
|
||||
assert_eq!(trimmed.as_ref(), &expected_items);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use arrow_array::{Array, RecordBatch};
|
||||
use arrow_data::ArrayData;
|
||||
|
||||
/// Counts memory used by buffers of Arrow arrays and RecordBatches.
|
||||
///
|
||||
/// This is meant to capture how much memory is being used by the Arrow data
|
||||
/// structures as they are. It does not represent the memory used if the data
|
||||
/// were to be serialized and then deserialized. In particular:
|
||||
///
|
||||
/// * This does not double count memory used by buffers shared by multiple
|
||||
/// arrays or batches. Round-tripped data may use more memory because of this.
|
||||
/// * This counts the **total** size of the buffers, even if the array is a slice.
|
||||
/// Round-tripped data may use less memory because of this.
|
||||
#[derive(Default)]
|
||||
pub struct MemoryAccumulator {
|
||||
seen: HashSet<usize>,
|
||||
total: usize,
|
||||
}
|
||||
|
||||
impl MemoryAccumulator {
|
||||
pub fn record_array(&mut self, array: &dyn Array) {
|
||||
let data = array.to_data();
|
||||
self.record_array_data(&data);
|
||||
}
|
||||
|
||||
fn record_array_data(&mut self, data: &ArrayData) {
|
||||
for buffer in data.buffers() {
|
||||
let ptr = buffer.as_ptr();
|
||||
if self.seen.insert(ptr as usize) {
|
||||
self.total += buffer.capacity();
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(nulls) = data.nulls() {
|
||||
let null_buf = nulls.inner().inner();
|
||||
let ptr = null_buf.as_ptr();
|
||||
if self.seen.insert(ptr as usize) {
|
||||
self.total += null_buf.capacity();
|
||||
}
|
||||
}
|
||||
|
||||
for child in data.child_data() {
|
||||
self.record_array_data(child);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn record_batch(&mut self, batch: &RecordBatch) {
|
||||
for array in batch.columns() {
|
||||
self.record_array(array);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn total(&self) -> usize {
|
||||
self.total
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::Int32Array;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_memory_accumulator() {
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
|
||||
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
|
||||
)
|
||||
.unwrap();
|
||||
let slice = batch.slice(1, 2);
|
||||
|
||||
let mut acc = MemoryAccumulator::default();
|
||||
|
||||
// Should record whole buffer, not just slice
|
||||
acc.record_batch(&slice);
|
||||
assert_eq!(acc.total(), 3 * std::mem::size_of::<i32>());
|
||||
|
||||
// Should not double count
|
||||
acc.record_batch(&slice);
|
||||
assert_eq!(acc.total(), 3 * std::mem::size_of::<i32>());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,304 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
use arrow_array::{ArrayRef, UInt64Array, make_array};
|
||||
use arrow_buffer::Buffer;
|
||||
use arrow_data::ArrayDataBuilder;
|
||||
use arrow_schema::{ArrowError, DataType};
|
||||
use arrow_select::take::take;
|
||||
|
||||
use crate::DataTypeExt;
|
||||
|
||||
type Result<T> = std::result::Result<T, ArrowError>;
|
||||
|
||||
pub const INLINE_VALUE_MAX_BYTES: usize = 32;
|
||||
|
||||
pub fn extract_scalar_value(array: &ArrayRef, idx: usize) -> Result<ArrayRef> {
|
||||
if idx >= array.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Scalar index out of bounds".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
take(array.as_ref(), &UInt64Array::from(vec![idx as u64]), None)
|
||||
}
|
||||
|
||||
fn read_u32(buf: &[u8], offset: &mut usize) -> Result<u32> {
|
||||
if *offset + 4 > buf.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid scalar value buffer: unexpected EOF".to_string(),
|
||||
));
|
||||
}
|
||||
let bytes = [
|
||||
buf[*offset],
|
||||
buf[*offset + 1],
|
||||
buf[*offset + 2],
|
||||
buf[*offset + 3],
|
||||
];
|
||||
*offset += 4;
|
||||
Ok(u32::from_le_bytes(bytes))
|
||||
}
|
||||
|
||||
fn read_bytes<'a>(buf: &'a [u8], offset: &mut usize, len: usize) -> Result<&'a [u8]> {
|
||||
if *offset + len > buf.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid scalar value buffer: unexpected EOF".to_string(),
|
||||
));
|
||||
}
|
||||
let slice = &buf[*offset..*offset + len];
|
||||
*offset += len;
|
||||
Ok(slice)
|
||||
}
|
||||
|
||||
fn write_u32(out: &mut Vec<u8>, v: u32) {
|
||||
out.extend_from_slice(&v.to_le_bytes());
|
||||
}
|
||||
|
||||
fn write_bytes(out: &mut Vec<u8>, bytes: &[u8]) {
|
||||
out.extend_from_slice(bytes);
|
||||
}
|
||||
|
||||
pub fn encode_scalar_value_buffer(scalar: &ArrayRef) -> Result<Vec<u8>> {
|
||||
if scalar.len() != 1 || scalar.null_count() != 0 {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Scalar value buffer must be a single non-null value".to_string(),
|
||||
));
|
||||
}
|
||||
let data = scalar.to_data();
|
||||
if data.offset() != 0 {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Scalar value buffer must have offset=0".to_string(),
|
||||
));
|
||||
}
|
||||
if !data.child_data().is_empty() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Scalar value buffer does not support nested types".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Minimal format (RFC): store the Arrow value buffers for a length-1 array.
|
||||
// Null bitmap and child data are intentionally not supported here.
|
||||
//
|
||||
// | u32 num_buffers |
|
||||
// | u32 buffer_0_len | ... | u32 buffer_{n-1}_len |
|
||||
// | buffer_0 bytes | ... | buffer_{n-1} bytes |
|
||||
let mut out = Vec::with_capacity(128);
|
||||
let buffers = data.buffers();
|
||||
write_u32(&mut out, buffers.len() as u32);
|
||||
for b in buffers {
|
||||
write_u32(&mut out, b.len() as u32);
|
||||
}
|
||||
for b in buffers {
|
||||
write_bytes(&mut out, b.as_slice());
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub fn decode_scalar_from_value_buffer(
|
||||
data_type: &DataType,
|
||||
value_buffer: &[u8],
|
||||
) -> Result<ArrayRef> {
|
||||
if matches!(
|
||||
data_type,
|
||||
DataType::Struct(_) | DataType::FixedSizeList(_, _)
|
||||
) {
|
||||
return Err(ArrowError::InvalidArgumentError(format!(
|
||||
"Scalar value buffer does not support nested data type {:?}",
|
||||
data_type
|
||||
)));
|
||||
}
|
||||
|
||||
let mut offset = 0;
|
||||
let num_buffers = read_u32(value_buffer, &mut offset)? as usize;
|
||||
let buffer_lens = (0..num_buffers)
|
||||
.map(|_| read_u32(value_buffer, &mut offset).map(|l| l as usize))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let mut buffers = Vec::with_capacity(num_buffers);
|
||||
for len in buffer_lens {
|
||||
let bytes = read_bytes(value_buffer, &mut offset, len)?;
|
||||
buffers.push(Buffer::from_vec(bytes.to_vec()));
|
||||
}
|
||||
|
||||
if offset != value_buffer.len() {
|
||||
return Err(ArrowError::InvalidArgumentError(
|
||||
"Invalid scalar value buffer: trailing bytes".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut builder = ArrayDataBuilder::new(data_type.clone())
|
||||
.len(1)
|
||||
.null_count(0);
|
||||
for b in buffers {
|
||||
builder = builder.add_buffer(b);
|
||||
}
|
||||
Ok(make_array(builder.build()?))
|
||||
}
|
||||
|
||||
pub fn decode_scalar_from_inline_value(
|
||||
data_type: &DataType,
|
||||
inline_value: &[u8],
|
||||
) -> Result<ArrayRef> {
|
||||
// I expect our input to be safe here, but I added some debug_assert_eq statements just in case.
|
||||
// If they are triggered, we may need to change them to return actual errors.
|
||||
//
|
||||
// Boolean values are bit-packed in Arrow and therefore are not "fixed-stride" in bytes.
|
||||
// As a result, `byte_width_opt()` returns `None` for `DataType::Boolean`, even though a
|
||||
// length-1 scalar can be represented inline using a single byte (matching `try_inline_value`).
|
||||
if matches!(data_type, DataType::Boolean) {
|
||||
debug_assert_eq!(
|
||||
inline_value.len(),
|
||||
1,
|
||||
"Invalid boolean inline scalar length (expected 1 byte, got {})",
|
||||
inline_value.len()
|
||||
);
|
||||
} else if let Some(byte_width) = data_type.byte_width_opt() {
|
||||
debug_assert_eq!(
|
||||
inline_value.len(),
|
||||
byte_width,
|
||||
"Inline constant length mismatch for {:?}: expected {} bytes but got {}",
|
||||
data_type,
|
||||
byte_width,
|
||||
inline_value.len()
|
||||
);
|
||||
}
|
||||
|
||||
let data = ArrayDataBuilder::new(data_type.clone())
|
||||
.len(1)
|
||||
.null_count(0)
|
||||
.add_buffer(Buffer::from_vec(inline_value.to_vec()))
|
||||
.build()?;
|
||||
Ok(make_array(data))
|
||||
}
|
||||
|
||||
pub fn try_inline_value(scalar: &ArrayRef) -> Option<Vec<u8>> {
|
||||
if scalar.null_count() != 0 || scalar.len() != 1 {
|
||||
return None;
|
||||
}
|
||||
let data = scalar.to_data();
|
||||
if !data.child_data().is_empty() {
|
||||
return None;
|
||||
}
|
||||
if data.buffers().len() != 1 {
|
||||
return None;
|
||||
}
|
||||
let bytes = data.buffers()[0].as_slice();
|
||||
if bytes.len() > INLINE_VALUE_MAX_BYTES {
|
||||
return None;
|
||||
}
|
||||
Some(bytes.to_vec())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{
|
||||
BooleanArray, DictionaryArray, FixedSizeBinaryArray, Int8Array, Int32Array, StringArray,
|
||||
cast::AsArray, types::Int8Type,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_extract_scalar_value() {
|
||||
let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]));
|
||||
let scalar = extract_scalar_value(&array, 2).unwrap();
|
||||
assert_eq!(scalar.len(), 1);
|
||||
assert_eq!(
|
||||
scalar
|
||||
.as_primitive::<arrow_array::types::Int32Type>()
|
||||
.value(0),
|
||||
3
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_scalar_value_from_full_dictionary() {
|
||||
let values = Arc::new(StringArray::from(
|
||||
(0..=i8::MAX)
|
||||
.map(|value| format!("value-{value}"))
|
||||
.collect::<Vec<_>>(),
|
||||
));
|
||||
let keys = Int8Array::from((0..=i8::MAX).collect::<Vec<_>>());
|
||||
let array: ArrayRef = Arc::new(DictionaryArray::<Int8Type>::new(keys, values));
|
||||
|
||||
let scalar = extract_scalar_value(&array, i8::MAX as usize).unwrap();
|
||||
|
||||
let scalar = scalar.as_dictionary::<Int8Type>();
|
||||
assert_eq!(scalar.len(), 1);
|
||||
assert_eq!(scalar.key(0), Some(i8::MAX as usize));
|
||||
assert_eq!(scalar.values().len(), i8::MAX as usize + 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scalar_value_buffer_utf8_round_trip() {
|
||||
let scalar: ArrayRef = Arc::new(StringArray::from(vec!["hello"]));
|
||||
let buf = encode_scalar_value_buffer(&scalar).unwrap();
|
||||
let decoded = decode_scalar_from_value_buffer(&DataType::Utf8, &buf).unwrap();
|
||||
assert_eq!(decoded.len(), 1);
|
||||
assert_eq!(decoded.null_count(), 0);
|
||||
assert_eq!(decoded.as_string::<i32>().value(0), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scalar_value_buffer_fixed_size_binary_round_trip() {
|
||||
let val = vec![0xABu8; 33];
|
||||
let scalar: ArrayRef = Arc::new(
|
||||
FixedSizeBinaryArray::try_from_sparse_iter_with_size(
|
||||
std::iter::once(Some(val.as_slice())),
|
||||
33,
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
let buf = encode_scalar_value_buffer(&scalar).unwrap();
|
||||
let decoded =
|
||||
decode_scalar_from_value_buffer(&DataType::FixedSizeBinary(33), &buf).unwrap();
|
||||
assert_eq!(decoded.len(), 1);
|
||||
assert_eq!(decoded.as_fixed_size_binary().value(0), val.as_slice());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inline_value_boolean_round_trip() {
|
||||
let scalar: ArrayRef = Arc::new(BooleanArray::from_iter([Some(true)]));
|
||||
let inline = try_inline_value(&scalar).unwrap();
|
||||
let decoded = decode_scalar_from_inline_value(&DataType::Boolean, &inline).unwrap();
|
||||
assert_eq!(decoded.len(), 1);
|
||||
assert_eq!(decoded.null_count(), 0);
|
||||
assert!(decoded.as_boolean().value(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scalar_value_buffer_rejects_nested_type() {
|
||||
let field = Arc::new(arrow_schema::Field::new("item", DataType::Int32, false));
|
||||
let list: ArrayRef = Arc::new(arrow_array::FixedSizeListArray::new(
|
||||
field,
|
||||
2,
|
||||
Arc::new(Int32Array::from(vec![1, 2])),
|
||||
None,
|
||||
));
|
||||
let scalar = list.slice(0, 1);
|
||||
assert!(encode_scalar_value_buffer(&scalar).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_scalar_from_value_buffer_rejects_nested_type() {
|
||||
let buf = Vec::<u8>::new();
|
||||
let res =
|
||||
decode_scalar_from_value_buffer(&DataType::Struct(arrow_schema::Fields::empty()), &buf);
|
||||
assert!(res.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_scalar_from_value_buffer_trailing_bytes() {
|
||||
// num_buffers = 0, plus an extra byte
|
||||
let mut bytes = Vec::new();
|
||||
bytes.extend_from_slice(&0u32.to_le_bytes());
|
||||
bytes.push(1);
|
||||
let res = decode_scalar_from_value_buffer(&DataType::Int32, &bytes);
|
||||
assert!(res.is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Extension to arrow schema
|
||||
|
||||
use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema};
|
||||
|
||||
use crate::{ARROW_EXT_NAME_KEY, BLOB_META_KEY, BLOB_V2_EXT_NAME};
|
||||
|
||||
pub enum Indentation {
|
||||
OneLine,
|
||||
MultiLine(u8),
|
||||
}
|
||||
|
||||
impl Indentation {
|
||||
fn value(&self) -> String {
|
||||
match self {
|
||||
Self::OneLine => "".to_string(),
|
||||
Self::MultiLine(spaces) => " ".repeat(*spaces as usize),
|
||||
}
|
||||
}
|
||||
|
||||
fn deepen(&self) -> Self {
|
||||
match self {
|
||||
Self::OneLine => Self::OneLine,
|
||||
Self::MultiLine(spaces) => Self::MultiLine(spaces + 2),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extends the functionality of [arrow_schema::Field].
|
||||
pub trait FieldExt {
|
||||
/// Create a compact string representation of the field
|
||||
///
|
||||
/// This is intended for display purposes and not for serialization
|
||||
fn to_compact_string(&self, indent: Indentation) -> String;
|
||||
|
||||
/// Check if the field is marked as a packed struct
|
||||
fn is_packed_struct(&self) -> bool;
|
||||
|
||||
/// Check if the field is marked as a blob
|
||||
fn is_blob(&self) -> bool;
|
||||
|
||||
/// Check if the field is marked as a blob
|
||||
fn is_blob_v2(&self) -> bool;
|
||||
}
|
||||
|
||||
impl FieldExt for Field {
|
||||
fn to_compact_string(&self, indent: Indentation) -> String {
|
||||
let mut result = format!("{}: ", self.name().clone());
|
||||
match self.data_type() {
|
||||
DataType::Struct(fields) => {
|
||||
result += "{";
|
||||
result += &indent.value();
|
||||
for (field_idx, field) in fields.iter().enumerate() {
|
||||
result += field.to_compact_string(indent.deepen()).as_str();
|
||||
if field_idx < fields.len() - 1 {
|
||||
result += ",";
|
||||
}
|
||||
result += indent.value().as_str();
|
||||
}
|
||||
result += "}";
|
||||
}
|
||||
DataType::List(field)
|
||||
| DataType::LargeList(field)
|
||||
| DataType::ListView(field)
|
||||
| DataType::LargeListView(field) => {
|
||||
result += "[";
|
||||
result += field.to_compact_string(indent.deepen()).as_str();
|
||||
result += "]";
|
||||
}
|
||||
DataType::FixedSizeList(child, dimension) => {
|
||||
result += &format!(
|
||||
"[{}; {}]",
|
||||
child.to_compact_string(indent.deepen()),
|
||||
dimension
|
||||
);
|
||||
}
|
||||
DataType::Dictionary(key_type, value_type) => {
|
||||
result += &value_type.to_string();
|
||||
result += "@";
|
||||
result += &key_type.to_string();
|
||||
}
|
||||
_ => {
|
||||
result += &self.data_type().to_string();
|
||||
}
|
||||
}
|
||||
if self.is_nullable() {
|
||||
result += "?";
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
// Check if field has metadata `packed` set to true, this check is case insensitive.
|
||||
fn is_packed_struct(&self) -> bool {
|
||||
let field_metadata = self.metadata();
|
||||
const PACKED_KEYS: [&str; 2] = ["packed", "lance-encoding:packed"];
|
||||
PACKED_KEYS.iter().any(|key| {
|
||||
field_metadata
|
||||
.get(*key)
|
||||
.map(|value| value.eq_ignore_ascii_case("true"))
|
||||
.unwrap_or(false)
|
||||
})
|
||||
}
|
||||
|
||||
fn is_blob(&self) -> bool {
|
||||
let field_metadata = self.metadata();
|
||||
field_metadata.get(BLOB_META_KEY).is_some()
|
||||
|| field_metadata
|
||||
.get(ARROW_EXT_NAME_KEY)
|
||||
.map(|value| value == BLOB_V2_EXT_NAME)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn is_blob_v2(&self) -> bool {
|
||||
let field_metadata = self.metadata();
|
||||
field_metadata
|
||||
.get(ARROW_EXT_NAME_KEY)
|
||||
.map(|value| value == BLOB_V2_EXT_NAME)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
|
||||
/// Extends the functionality of [arrow_schema::Schema].
|
||||
pub trait SchemaExt {
|
||||
/// Create a new [`Schema`] with one extra field.
|
||||
fn try_with_column(&self, field: Field) -> std::result::Result<Schema, ArrowError>;
|
||||
|
||||
fn try_with_column_at(
|
||||
&self,
|
||||
index: usize,
|
||||
field: Field,
|
||||
) -> std::result::Result<Schema, ArrowError>;
|
||||
|
||||
fn field_names(&self) -> Vec<&String>;
|
||||
|
||||
fn without_column(&self, column_name: &str) -> Schema;
|
||||
|
||||
/// Create a compact string representation of the schema
|
||||
///
|
||||
/// This is intended for display purposes and not for serialization
|
||||
fn to_compact_string(&self, indent: Indentation) -> String;
|
||||
}
|
||||
|
||||
impl SchemaExt for Schema {
|
||||
fn try_with_column(&self, field: Field) -> std::result::Result<Schema, ArrowError> {
|
||||
if self.column_with_name(field.name()).is_some() {
|
||||
return Err(ArrowError::SchemaError(format!(
|
||||
"Can not append column {} on schema: {:?}",
|
||||
field.name(),
|
||||
self
|
||||
)));
|
||||
};
|
||||
let mut fields: Vec<FieldRef> = self.fields().iter().cloned().collect();
|
||||
fields.push(FieldRef::new(field));
|
||||
Ok(Self::new_with_metadata(fields, self.metadata.clone()))
|
||||
}
|
||||
|
||||
fn try_with_column_at(
|
||||
&self,
|
||||
index: usize,
|
||||
field: Field,
|
||||
) -> std::result::Result<Schema, ArrowError> {
|
||||
if self.column_with_name(field.name()).is_some() {
|
||||
return Err(ArrowError::SchemaError(format!(
|
||||
"Failed to modify schema: Inserting column {} would create a duplicate column in schema: {:?}",
|
||||
field.name(),
|
||||
self
|
||||
)));
|
||||
};
|
||||
let mut fields: Vec<FieldRef> = self.fields().iter().cloned().collect();
|
||||
fields.insert(index, FieldRef::new(field));
|
||||
Ok(Self::new_with_metadata(fields, self.metadata.clone()))
|
||||
}
|
||||
|
||||
/// Project the schema to remove the given column.
|
||||
///
|
||||
/// This only works on top-level fields right now. If a field does not exist,
|
||||
/// the schema will be returned as is.
|
||||
fn without_column(&self, column_name: &str) -> Schema {
|
||||
let fields: Vec<FieldRef> = self
|
||||
.fields()
|
||||
.iter()
|
||||
.filter(|f| f.name() != column_name)
|
||||
.cloned()
|
||||
.collect();
|
||||
Self::new_with_metadata(fields, self.metadata.clone())
|
||||
}
|
||||
|
||||
fn field_names(&self) -> Vec<&String> {
|
||||
self.fields().iter().map(|f| f.name()).collect()
|
||||
}
|
||||
|
||||
fn to_compact_string(&self, indent: Indentation) -> String {
|
||||
let mut result = "{".to_string();
|
||||
result += &indent.value();
|
||||
for (field_idx, field) in self.fields.iter().enumerate() {
|
||||
result += field.to_compact_string(indent.deepen()).as_str();
|
||||
if field_idx < self.fields.len() - 1 {
|
||||
result += ",";
|
||||
}
|
||||
result += indent.value().as_str();
|
||||
}
|
||||
result += "}";
|
||||
result
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,507 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Utilities for working with streams of [`RecordBatch`].
|
||||
|
||||
use arrow_array::RecordBatch;
|
||||
use arrow_schema::{ArrowError, SchemaRef};
|
||||
use futures::stream::{self, Stream, StreamExt};
|
||||
use std::pin::Pin;
|
||||
|
||||
use crate::deepcopy::deep_copy_batch_sliced;
|
||||
|
||||
/// Rechunks a stream of [`RecordBatch`] so that each output batch has
|
||||
/// approximately `target_bytes` of array data.
|
||||
///
|
||||
/// Small input batches are accumulated (by concatenation) until at least
|
||||
/// `min_bytes` of data has been collected. If the resulting batch exceeds
|
||||
/// `max_bytes`, it is sliced into roughly equal pieces of ~`max_bytes`
|
||||
/// (assuming uniform row sizes).
|
||||
pub fn rechunk_stream_by_size<S, E>(
|
||||
input: S,
|
||||
input_schema: SchemaRef,
|
||||
min_bytes: usize,
|
||||
max_bytes: usize,
|
||||
) -> impl Stream<Item = Result<RecordBatch, E>>
|
||||
where
|
||||
S: Stream<Item = Result<RecordBatch, E>>,
|
||||
E: From<ArrowError>,
|
||||
{
|
||||
rechunk_stream_by_size_inner(input, input_schema, min_bytes, max_bytes, false)
|
||||
}
|
||||
|
||||
/// Like [`rechunk_stream_by_size`] but deep-copies slices so that
|
||||
/// `get_array_memory_size` reflects the true size of each output batch.
|
||||
///
|
||||
/// After a normal `RecordBatch::slice`, the backing buffers are shared with
|
||||
/// the original batch, so `get_array_memory_size` still reports the full
|
||||
/// parent size. This variant deep-copies every slice produced during the
|
||||
/// splitting phase, which allows the stream to detect and re-split slices
|
||||
/// that still exceed `max_bytes` (e.g. because a single row is much larger
|
||||
/// than average).
|
||||
///
|
||||
/// The deep copy is a last resort and potentially expensive for large
|
||||
/// batches. However, it is only performed when a batch actually needs to be
|
||||
/// sliced — batches that are already within the target range pass through at
|
||||
/// zero cost. Use this only when the hard cap on `max_bytes` is a
|
||||
/// correctness requirement, not merely a performance hint.
|
||||
pub fn rechunk_stream_by_size_deep_copy<S, E>(
|
||||
input: S,
|
||||
input_schema: SchemaRef,
|
||||
min_bytes: usize,
|
||||
max_bytes: usize,
|
||||
) -> impl Stream<Item = Result<RecordBatch, E>>
|
||||
where
|
||||
S: Stream<Item = Result<RecordBatch, E>>,
|
||||
E: From<ArrowError>,
|
||||
{
|
||||
rechunk_stream_by_size_inner(input, input_schema, min_bytes, max_bytes, true)
|
||||
}
|
||||
|
||||
fn rechunk_stream_by_size_inner<S, E>(
|
||||
input: S,
|
||||
input_schema: SchemaRef,
|
||||
min_bytes: usize,
|
||||
max_bytes: usize,
|
||||
deep_copy: bool,
|
||||
) -> impl Stream<Item = Result<RecordBatch, E>>
|
||||
where
|
||||
S: Stream<Item = Result<RecordBatch, E>>,
|
||||
E: From<ArrowError>,
|
||||
{
|
||||
stream::try_unfold(
|
||||
RechunkState {
|
||||
input: Box::pin(input),
|
||||
accumulated: Vec::new(),
|
||||
acc_bytes: 0,
|
||||
done: false,
|
||||
input_schema,
|
||||
min_bytes,
|
||||
max_bytes,
|
||||
deep_copy,
|
||||
},
|
||||
|mut state| async move {
|
||||
if state.done && state.accumulated.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Pull batches until we reach the byte target or exhaust input.
|
||||
// Always pull at least one batch so that min_bytes=0 works.
|
||||
while !state.done && (state.accumulated.is_empty() || state.acc_bytes < state.min_bytes)
|
||||
{
|
||||
match state.input.next().await {
|
||||
Some(Ok(batch)) => {
|
||||
state.acc_bytes += batch.get_array_memory_size();
|
||||
state.accumulated.push(batch);
|
||||
}
|
||||
Some(Err(e)) => return Err(e),
|
||||
None => {
|
||||
state.done = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if state.accumulated.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Fast path: if the first accumulated batch already meets the
|
||||
// byte threshold, deliver it directly instead of concatenating
|
||||
// everything together (which would just get sliced back apart).
|
||||
if state.accumulated.len() > 1
|
||||
&& state.accumulated[0].get_array_memory_size() >= state.min_bytes
|
||||
{
|
||||
let b = state.accumulated.remove(0);
|
||||
state.acc_bytes -= b.get_array_memory_size();
|
||||
return Ok(Some((b, state)));
|
||||
}
|
||||
|
||||
let batch = if state.accumulated.len() == 1 {
|
||||
state.accumulated.pop().unwrap()
|
||||
} else {
|
||||
let b =
|
||||
arrow_select::concat::concat_batches(&state.input_schema, &state.accumulated)
|
||||
.map_err(E::from)?;
|
||||
state.accumulated.clear();
|
||||
b
|
||||
};
|
||||
state.acc_bytes = 0;
|
||||
|
||||
// Slice the batch into ~max_bytes pieces assuming uniform row sizes.
|
||||
let mut slices =
|
||||
slice_batch(batch, state.max_bytes, state.deep_copy).map_err(E::from)?;
|
||||
|
||||
if slices.len() == 1 {
|
||||
Ok(Some((slices.pop().unwrap(), state)))
|
||||
} else {
|
||||
let first = slices.remove(0);
|
||||
|
||||
// Stash leftover slices for subsequent iterations.
|
||||
for a in &slices {
|
||||
state.acc_bytes += a.get_array_memory_size();
|
||||
}
|
||||
state.accumulated = slices;
|
||||
|
||||
Ok(Some((first, state)))
|
||||
}
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Slice a batch into pieces of at most `max_bytes`.
|
||||
///
|
||||
/// When `deep_copy` is false, slices share buffers with the original batch
|
||||
/// and `get_array_memory_size` will still report the parent buffer size.
|
||||
/// This is fine when the caller only needs approximate sizing.
|
||||
///
|
||||
/// When `deep_copy` is true, each slice is deep-copied so that
|
||||
/// `get_array_memory_size` reflects the true size. If a deep-copied slice
|
||||
/// still exceeds `max_bytes` (due to non-uniform row sizes), it is
|
||||
/// recursively split until every piece is within budget or contains only a
|
||||
/// single row.
|
||||
fn slice_batch(
|
||||
batch: RecordBatch,
|
||||
max_bytes: usize,
|
||||
deep_copy: bool,
|
||||
) -> Result<Vec<RecordBatch>, ArrowError> {
|
||||
let batch_bytes = batch.get_array_memory_size();
|
||||
let num_rows = batch.num_rows();
|
||||
|
||||
if batch_bytes <= max_bytes || num_rows <= 1 {
|
||||
return Ok(vec![batch]);
|
||||
}
|
||||
|
||||
let rows_per_chunk = (max_bytes as u64 * num_rows as u64 / batch_bytes as u64).max(1) as usize;
|
||||
|
||||
let mut result = Vec::new();
|
||||
let mut offset = 0;
|
||||
while offset < num_rows {
|
||||
let len = rows_per_chunk.min(num_rows - offset);
|
||||
let slice = batch.slice(offset, len);
|
||||
if deep_copy {
|
||||
let copied = deep_copy_batch_sliced(&slice)?;
|
||||
// Recurse: the deep-copied slice has accurate sizes, so if it
|
||||
// still exceeds max_bytes we can split further.
|
||||
result.extend(slice_batch(copied, max_bytes, true)?);
|
||||
} else {
|
||||
result.push(slice);
|
||||
}
|
||||
offset += len;
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Internal state for [`rechunk_stream`].
|
||||
///
|
||||
/// Kept as a named struct so the `try_unfold` closure stays readable.
|
||||
struct RechunkState<S> {
|
||||
input: Pin<Box<S>>,
|
||||
accumulated: Vec<RecordBatch>,
|
||||
acc_bytes: usize,
|
||||
done: bool,
|
||||
input_schema: SchemaRef,
|
||||
min_bytes: usize,
|
||||
max_bytes: usize,
|
||||
deep_copy: bool,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::Int32Array;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use futures::executor::block_on;
|
||||
|
||||
fn make_batch(num_rows: usize) -> RecordBatch {
|
||||
let schema = test_schema();
|
||||
let values: Vec<i32> = (0..num_rows as i32).collect();
|
||||
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))]).unwrap()
|
||||
}
|
||||
|
||||
fn test_schema() -> SchemaRef {
|
||||
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
|
||||
}
|
||||
|
||||
fn collect_rechunked(
|
||||
batches: Vec<RecordBatch>,
|
||||
min_bytes: usize,
|
||||
max_bytes: usize,
|
||||
) -> Vec<RecordBatch> {
|
||||
let input = stream::iter(batches.into_iter().map(Ok::<_, ArrowError>));
|
||||
let rechunked = rechunk_stream_by_size(input, test_schema(), min_bytes, max_bytes);
|
||||
block_on(rechunked.collect::<Vec<_>>())
|
||||
.into_iter()
|
||||
.map(|r| r.unwrap())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn total_rows(batches: &[RecordBatch]) -> usize {
|
||||
batches.iter().map(|b| b.num_rows()).sum()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_stream() {
|
||||
let result = collect_rechunked(vec![], 100, 200);
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_single_batch_passthrough() {
|
||||
let batch = make_batch(100);
|
||||
let bytes = batch.get_array_memory_size();
|
||||
// Batch is between min and max — should pass through as-is.
|
||||
let result = collect_rechunked(vec![batch], bytes / 2, bytes * 2);
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].num_rows(), 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_small_batches_concatenated() {
|
||||
let one_batch_bytes = make_batch(10).get_array_memory_size();
|
||||
let batches: Vec<_> = (0..8).map(|_| make_batch(10)).collect();
|
||||
// min = 5 batches worth, max = 10 batches worth.
|
||||
let result = collect_rechunked(batches, one_batch_bytes * 5, one_batch_bytes * 10);
|
||||
assert_eq!(total_rows(&result), 80);
|
||||
// Should have been concatenated into fewer batches than the 8 inputs.
|
||||
assert!(
|
||||
result.len() < 8,
|
||||
"expected fewer output batches, got {}",
|
||||
result.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_large_batch_sliced() {
|
||||
let batch = make_batch(1000);
|
||||
let bytes = batch.get_array_memory_size();
|
||||
let result = collect_rechunked(vec![batch], bytes / 8, bytes / 4);
|
||||
assert_eq!(total_rows(&result), 1000);
|
||||
assert!(
|
||||
result.len() >= 4,
|
||||
"expected at least 4 slices, got {}",
|
||||
result.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sliced_leftovers_are_not_recombined() {
|
||||
// Key test for the fast-path optimisation. When a large batch is
|
||||
// sliced, leftover slices should be delivered one-at-a-time without
|
||||
// being concatenated back together. We verify this by checking that
|
||||
// every output buffer pointer falls inside the original batch's
|
||||
// allocation (i.e. they are all zero-copy slices, not fresh copies).
|
||||
let batch = make_batch(1000);
|
||||
let bytes = batch.get_array_memory_size();
|
||||
let orig_data = batch.column(0).to_data();
|
||||
let orig_buf = &orig_data.buffers()[0];
|
||||
let orig_start = orig_buf.as_ptr() as usize;
|
||||
let orig_end = orig_start + orig_buf.len();
|
||||
|
||||
let result = collect_rechunked(vec![batch], bytes / 8, bytes / 4);
|
||||
|
||||
assert_eq!(total_rows(&result), 1000);
|
||||
assert!(result.len() >= 4);
|
||||
|
||||
for (i, b) in result.iter().enumerate() {
|
||||
let ptr = b.column(0).to_data().buffers()[0].as_ptr() as usize;
|
||||
assert!(
|
||||
ptr >= orig_start && ptr < orig_end,
|
||||
"slice {i} buffer at {ptr:#x} is outside the original allocation \
|
||||
[{orig_start:#x}, {orig_end:#x}) — it was re-concatenated"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flush_remainder_on_stream_end() {
|
||||
// Data below min_bytes should still be flushed when the stream ends.
|
||||
let batch = make_batch(10);
|
||||
let bytes = batch.get_array_memory_size();
|
||||
let result = collect_rechunked(vec![batch], bytes * 100, bytes * 200);
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].num_rows(), 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_large_then_small_batches() {
|
||||
// After a large batch is fully drained, subsequent small batches
|
||||
// should be accumulated normally.
|
||||
let large = make_batch(1000);
|
||||
let small_bytes = make_batch(10).get_array_memory_size();
|
||||
let batches = vec![
|
||||
large,
|
||||
make_batch(10),
|
||||
make_batch(10),
|
||||
make_batch(10),
|
||||
make_batch(10),
|
||||
make_batch(10),
|
||||
];
|
||||
let result = collect_rechunked(batches, small_bytes * 3, small_bytes * 100);
|
||||
assert_eq!(total_rows(&result), 1050);
|
||||
// The large batch should appear (possibly sliced) followed by
|
||||
// concatenated small batches, so we should have fewer output batches
|
||||
// than the 6 inputs.
|
||||
assert!(result.len() < 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_row_preservation_across_slicing() {
|
||||
// Verify that every input row appears exactly once in the output
|
||||
// and in the correct order after slicing.
|
||||
let batch = make_batch(237); // odd count to exercise remainder slice
|
||||
let bytes = batch.get_array_memory_size();
|
||||
let result = collect_rechunked(vec![batch], bytes / 8, bytes / 5);
|
||||
|
||||
assert_eq!(total_rows(&result), 237);
|
||||
|
||||
let values: Vec<i32> = result
|
||||
.iter()
|
||||
.flat_map(|b| {
|
||||
b.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<Int32Array>()
|
||||
.unwrap()
|
||||
.values()
|
||||
.iter()
|
||||
.copied()
|
||||
})
|
||||
.collect();
|
||||
let expected: Vec<i32> = (0..237).collect();
|
||||
assert_eq!(values, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_min_bytes_zero_still_yields_all_rows() {
|
||||
// When min_bytes=0, the stream should still yield every batch.
|
||||
// This is the "chop only, don't coalesce" use case.
|
||||
let batches: Vec<_> = (0..5).map(|_| make_batch(100)).collect();
|
||||
let batch_bytes = batches[0].get_array_memory_size();
|
||||
let result = collect_rechunked(batches, 0, batch_bytes * 2);
|
||||
assert_eq!(total_rows(&result), 500);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_min_bytes_zero_slices_oversized() {
|
||||
// min_bytes=0 with a small max_bytes should still slice large batches.
|
||||
let batch = make_batch(1000);
|
||||
let bytes = batch.get_array_memory_size();
|
||||
let result = collect_rechunked(vec![batch], 0, bytes / 4);
|
||||
assert_eq!(total_rows(&result), 1000);
|
||||
assert!(
|
||||
result.len() >= 4,
|
||||
"expected at least 4 slices, got {}",
|
||||
result.len()
|
||||
);
|
||||
}
|
||||
|
||||
/// Build a batch with one variable-length string column.
|
||||
/// Every row is `small_size` bytes except the row at index `big_row_idx`
|
||||
/// which is `big_size` bytes.
|
||||
fn make_variable_batch(
|
||||
num_rows: usize,
|
||||
small_size: usize,
|
||||
big_row_idx: usize,
|
||||
big_size: usize,
|
||||
) -> RecordBatch {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
|
||||
let values: Vec<String> = (0..num_rows)
|
||||
.map(|i| {
|
||||
if i == big_row_idx {
|
||||
"X".repeat(big_size)
|
||||
} else {
|
||||
"x".repeat(small_size)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let array = arrow_array::StringArray::from(values);
|
||||
RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
|
||||
}
|
||||
|
||||
fn variable_schema() -> SchemaRef {
|
||||
Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]))
|
||||
}
|
||||
|
||||
fn collect_rechunked_variable(
|
||||
batches: Vec<RecordBatch>,
|
||||
min_bytes: usize,
|
||||
max_bytes: usize,
|
||||
) -> Vec<RecordBatch> {
|
||||
let input = stream::iter(batches.into_iter().map(Ok::<_, ArrowError>));
|
||||
let rechunked =
|
||||
rechunk_stream_by_size_deep_copy(input, variable_schema(), min_bytes, max_bytes);
|
||||
block_on(rechunked.collect::<Vec<_>>())
|
||||
.into_iter()
|
||||
.map(|r| r.unwrap())
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oversized_row_at_end() {
|
||||
// 100 rows: 99 small (64 bytes each) + 1 large (100KiB) at the end.
|
||||
let batch = make_variable_batch(100, 64, 99, 100 * 1024);
|
||||
let max_bytes = 64 * 1024;
|
||||
let result = collect_rechunked_variable(vec![batch], 0, max_bytes);
|
||||
assert_eq!(total_rows(&result), 100);
|
||||
for (i, b) in result.iter().enumerate() {
|
||||
let size = b.get_array_memory_size();
|
||||
assert!(
|
||||
size <= max_bytes || b.num_rows() == 1,
|
||||
"batch {i} has {size} bytes (max {max_bytes}) and {} rows",
|
||||
b.num_rows()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oversized_row_at_start() {
|
||||
// 100 rows: 1 large (100KiB) at the start + 99 small (64 bytes each).
|
||||
let batch = make_variable_batch(100, 64, 0, 100 * 1024);
|
||||
let max_bytes = 64 * 1024;
|
||||
let result = collect_rechunked_variable(vec![batch], 0, max_bytes);
|
||||
assert_eq!(total_rows(&result), 100);
|
||||
for (i, b) in result.iter().enumerate() {
|
||||
let size = b.get_array_memory_size();
|
||||
assert!(
|
||||
size <= max_bytes || b.num_rows() == 1,
|
||||
"batch {i} has {size} bytes (max {max_bytes}) and {} rows",
|
||||
b.num_rows()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oversized_row_in_middle() {
|
||||
// 100 rows: 1 large (100KiB) in the middle + 99 small (64 bytes each).
|
||||
let batch = make_variable_batch(100, 64, 50, 100 * 1024);
|
||||
let max_bytes = 64 * 1024;
|
||||
let result = collect_rechunked_variable(vec![batch], 0, max_bytes);
|
||||
assert_eq!(total_rows(&result), 100);
|
||||
for (i, b) in result.iter().enumerate() {
|
||||
let size = b.get_array_memory_size();
|
||||
assert!(
|
||||
size <= max_bytes || b.num_rows() == 1,
|
||||
"batch {i} has {size} bytes (max {max_bytes}) and {} rows",
|
||||
b.num_rows()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_error_propagation() {
|
||||
let input = stream::iter(vec![
|
||||
Ok(make_batch(10)),
|
||||
Err(ArrowError::ComputeError("boom".into())),
|
||||
Ok(make_batch(10)),
|
||||
]);
|
||||
let rechunked = rechunk_stream_by_size(input, test_schema(), 1, usize::MAX);
|
||||
let results: Vec<Result<RecordBatch, ArrowError>> = block_on(rechunked.collect());
|
||||
assert!(results.iter().any(|r| r.is_err()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The Lance Authors
|
||||
|
||||
//! Extension to arrow struct arrays
|
||||
|
||||
use arrow_array::{Array, StructArray, cast::AsArray, make_array};
|
||||
use arrow_buffer::NullBuffer;
|
||||
use arrow_data::{ArrayData, ArrayDataBuilder};
|
||||
use arrow_schema::ArrowError;
|
||||
|
||||
pub trait StructArrayExt {
|
||||
/// Removes the offset / length of the struct array by pushing it into the children
|
||||
///
|
||||
/// In arrow-rs when slice is called it recursively slices the children.
|
||||
/// In arrow-cpp when slice is called it just sets the offset/length of
|
||||
/// the struct array and leaves the children as-is
|
||||
///
|
||||
/// Both are legal approaches (╥﹏╥)
|
||||
///
|
||||
/// This method helps reduce complexity by folding into the arrow-rs approach
|
||||
fn normalize_slicing(&self) -> Result<Self, ArrowError>
|
||||
where
|
||||
Self: Sized;
|
||||
|
||||
/// Structs are allowed to mask valid items. For example, a struct array might be:
|
||||
///
|
||||
/// [ {"items": [1, 2, 3]}, NULL, {"items": NULL}]
|
||||
///
|
||||
/// However, the underlying items array might be: [[1, 2, 3], [4, 5], NULL]
|
||||
///
|
||||
/// The [4, 5] list is masked out because the struct array is null.
|
||||
///
|
||||
/// The struct validity would be [true, false, true] and the list validity would be [true, true, false]
|
||||
///
|
||||
/// This method pushes nulls down into all children. In the above example the list validity would become
|
||||
/// [true, false, false].
|
||||
///
|
||||
/// This method is not recursive. If a child is a struct array it will not push that child's nulls down.
|
||||
///
|
||||
/// This method does not remove garbage lists. It only updates the validity so a future call to
|
||||
/// [crate::list::ListArrayExt::filter_garbage_nulls] will remove the garbage lists (without
|
||||
/// this pushdown it would not)
|
||||
fn pushdown_nulls(&self) -> Result<Self, ArrowError>
|
||||
where
|
||||
Self: Sized;
|
||||
}
|
||||
|
||||
fn normalized_struct_array_data(data: ArrayData) -> Result<ArrayData, ArrowError> {
|
||||
let parent_offset = data.offset();
|
||||
let parent_len = data.len();
|
||||
let modified_children = data
|
||||
.child_data()
|
||||
.iter()
|
||||
.map(|d| {
|
||||
let d = normalized_struct_array_data(d.clone())?;
|
||||
let offset = d.offset();
|
||||
let len = d.len();
|
||||
if len < parent_len + parent_offset {
|
||||
return Err(ArrowError::InvalidArgumentError(format!(
|
||||
"Child array {} has length {} which is less than the parent length {} plus the parent offset {}",
|
||||
d.data_type(),
|
||||
len,
|
||||
parent_len,
|
||||
parent_offset
|
||||
)));
|
||||
}
|
||||
let new_offset = offset + parent_offset;
|
||||
d.into_builder().offset(new_offset)
|
||||
.len(parent_len)
|
||||
.build()
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
ArrayDataBuilder::new(data.data_type().clone())
|
||||
.len(parent_len)
|
||||
.offset(0)
|
||||
.buffers(data.buffers().to_vec())
|
||||
.child_data(modified_children)
|
||||
.build()
|
||||
}
|
||||
|
||||
impl StructArrayExt for StructArray {
|
||||
fn normalize_slicing(&self) -> Result<Self, ArrowError>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
if self.offset() == 0 && self.columns().iter().all(|c| c.len() == self.len()) {
|
||||
return Ok(self.clone());
|
||||
}
|
||||
|
||||
let data = normalized_struct_array_data(self.to_data())?;
|
||||
Ok(Self::from(data))
|
||||
}
|
||||
|
||||
fn pushdown_nulls(&self) -> Result<Self, ArrowError>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
let Some(validity) = self.nulls() else {
|
||||
return Ok(self.clone());
|
||||
};
|
||||
let data = self.to_data();
|
||||
let children = data
|
||||
.child_data()
|
||||
.iter()
|
||||
.map(|c| {
|
||||
if let Some(child_validity) = c.nulls() {
|
||||
let new_validity = child_validity.inner() & validity.inner();
|
||||
c.clone()
|
||||
.into_builder()
|
||||
.nulls(Some(NullBuffer::from(new_validity)))
|
||||
.build()
|
||||
} else {
|
||||
Ok(c.clone()
|
||||
.into_builder()
|
||||
.nulls(Some(validity.clone()))
|
||||
.build()?)
|
||||
}
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let arr = make_array(data.into_builder().child_data(children).build()?);
|
||||
Ok(arr.as_struct().clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use arrow_array::{Array, Int32Array, StructArray, cast::AsArray, make_array};
|
||||
use arrow_schema::{DataType, Field, Fields};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::r#struct::StructArrayExt;
|
||||
|
||||
#[test]
|
||||
fn test_normalize_slicing_no_offset() {
|
||||
let x = Int32Array::from(vec![1, 2, 3]);
|
||||
let y = Int32Array::from(vec![4, 5, 6]);
|
||||
let struct_array = StructArray::new(
|
||||
Fields::from(vec![
|
||||
Field::new("x", DataType::Int32, true),
|
||||
Field::new("y", DataType::Int32, true),
|
||||
]),
|
||||
vec![Arc::new(x), Arc::new(y)],
|
||||
None,
|
||||
);
|
||||
|
||||
let normalized = struct_array.normalize_slicing().unwrap();
|
||||
assert_eq!(normalized, struct_array);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_arrow_rs_slicing() {
|
||||
let x = Int32Array::from(vec![1, 2, 3, 4]);
|
||||
let y = Int32Array::from(vec![5, 6, 7, 8]);
|
||||
let struct_array = StructArray::new(
|
||||
Fields::from(vec![
|
||||
Field::new("x", DataType::Int32, true),
|
||||
Field::new("y", DataType::Int32, true),
|
||||
]),
|
||||
vec![Arc::new(x), Arc::new(y)],
|
||||
None,
|
||||
);
|
||||
|
||||
// Slicing with arrow-rs propagates the slicing to the children so there should
|
||||
// be no change needed to the struct array
|
||||
let sliced = struct_array.slice(1, 2);
|
||||
let normalized = sliced.normalize_slicing().unwrap();
|
||||
|
||||
assert_eq!(normalized, sliced);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_arrow_cpp_slicing() {
|
||||
let x = Int32Array::from(vec![1, 2, 3, 4]);
|
||||
let y = Int32Array::from(vec![5, 6, 7, 8]);
|
||||
let struct_array = StructArray::new(
|
||||
Fields::from(vec![
|
||||
Field::new("x", DataType::Int32, true),
|
||||
Field::new("y", DataType::Int32, true),
|
||||
]),
|
||||
vec![Arc::new(x), Arc::new(y)],
|
||||
None,
|
||||
);
|
||||
|
||||
let data = struct_array.to_data();
|
||||
let sliced = data.into_builder().offset(1).len(2).build().unwrap();
|
||||
let sliced = make_array(sliced);
|
||||
let normalized = sliced.as_struct().clone().normalize_slicing().unwrap();
|
||||
|
||||
assert_eq!(normalized, struct_array.slice(1, 2));
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user