mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-28 08:58:41 +00:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 05fce57841 | |||
| c94d9a2a16 | |||
| 6c8aa22704 | |||
| 84f46df876 | |||
| 43465f06d5 | |||
| bc762926ab | |||
| 2426f275cd |
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.38.0-beta.11"
|
||||
current_version = "0.38.0-beta.12"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -40,40 +40,31 @@ jobs:
|
||||
- target: aarch64-apple-darwin
|
||||
host: macos-latest
|
||||
features: fp16kernels
|
||||
# Fat LTO was ~111 of this job's ~113 minutes.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
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
|
||||
features: ","
|
||||
# The lower peak also keeps this on the standard 4-core runner.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc ninja nasm
|
||||
tail -n 1000 /c/ProgramData/chocolatey/logs/chocolatey.log
|
||||
# There is an issue where choco doesn't add nasm to the path
|
||||
export PATH="$PATH:/c/Program Files/NASM"
|
||||
nasm -v
|
||||
# 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
|
||||
features: ","
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc
|
||||
rustup target add aarch64-pc-windows-msvc
|
||||
# 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
|
||||
host: ubuntu-latest
|
||||
features: fp16kernels
|
||||
@@ -103,6 +94,14 @@ jobs:
|
||||
# https://github.com/napi-rs/napi-rs/blob/main/debian-aarch64.Dockerfile
|
||||
docker: ghcr.io/napi-rs/napi-rs/nodejs-rust:lts-debian-aarch64
|
||||
features: "fp16kernels"
|
||||
# Fat LTO OOM-killed rustc every nightly; even with lld it peaked
|
||||
# at 31391 MiB of the runner's 32 GiB.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
# arm64 Linux links through GNU `ld` where x86_64 defaults to
|
||||
# `rust-lld`, which is why only arm64 OOM'd. lld cut the largest
|
||||
# linker process 7.0 -> 4.0 GiB (lancedb/sophon#7313).
|
||||
linker: /tmp/aarch64-lld-clang
|
||||
pre_build: |-
|
||||
set -e &&
|
||||
apt-get update &&
|
||||
@@ -112,9 +111,30 @@ jobs:
|
||||
# AT_HWCAP2 (added in Linux 3.17). Define it for aws-lc-sys.
|
||||
export CFLAGS="$CFLAGS -DAT_HWCAP2=26" &&
|
||||
rustup target add aarch64-unknown-linux-gnu
|
||||
# Not `&&`-chained: in dash, errexit does not fire for a
|
||||
# non-final command in an `&&` list, so failures were ignored.
|
||||
#
|
||||
# A wrapper rather than `-C link-arg` because the per-target
|
||||
# rustflags variable does not reach every unit that links, while
|
||||
# the linker variable does. `clang` because GCC silently ignores
|
||||
# `-fuse-ld=lld` unless built with lld support. Two echoes
|
||||
# because printf's newline escape gets rewritten to `;` between
|
||||
# here and the container.
|
||||
echo '#!/bin/sh' > /tmp/aarch64-lld-clang
|
||||
echo 'exec clang --target=aarch64-unknown-linux-gnu --sysroot=/usr/aarch64-unknown-linux-gnu/aarch64-unknown-linux-gnu/sysroot --gcc-toolchain=/usr/aarch64-unknown-linux-gnu -fuse-ld=lld "$@"' >> /tmp/aarch64-lld-clang
|
||||
chmod 0755 /tmp/aarch64-lld-clang
|
||||
# Fail now, not at the cdylib link ~30 minutes later. Linking at
|
||||
# all also proves lld resolved; clang errors out when it cannot.
|
||||
echo 'int main(void){return 0;}' > /tmp/probe.c
|
||||
/tmp/aarch64-lld-clang /tmp/probe.c -o /tmp/probe
|
||||
readelf -h /tmp/probe | grep AArch64
|
||||
- target: aarch64-unknown-linux-musl
|
||||
host: ubuntu-2404-8x-x64
|
||||
features: ","
|
||||
# Fat LTO took the whole runner down. lld cannot help: it died
|
||||
# inside rustc's LLVM, before any linker was spawned.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
pre_build: |-
|
||||
set -e &&
|
||||
sudo apt-get update &&
|
||||
@@ -123,6 +143,19 @@ jobs:
|
||||
export EXTRA_ARGS="-x"
|
||||
name: build - ${{ matrix.settings.target }}
|
||||
runs-on: ${{ matrix.settings.host }}
|
||||
# On the job, not exported from `pre_build`: `Swatinem/rust-cache` hashes
|
||||
# `CARGO_*` into its cache key before any step runs, so a step-local export
|
||||
# leaves the key unchanged while cargo still rebuilds cold. The ThinLTO
|
||||
# legs had been doing that every run.
|
||||
#
|
||||
# Not `RUSTFLAGS`: setting it, even to "", discards every config-file
|
||||
# rustflag, silently dropping .cargo/config.toml's `target-cpu` and
|
||||
# `target-feature` from the published binaries.
|
||||
env:
|
||||
CARGO_PROFILE_RELEASE_LTO: ${{ matrix.settings.lto || 'fat' }}
|
||||
CARGO_PROFILE_RELEASE_CODEGEN_UNITS: ${{ matrix.settings.codegen_units || '1' }}
|
||||
# Empty elsewhere: a per-target variable is only read for that triple.
|
||||
CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER: ${{ matrix.settings.linker }}
|
||||
defaults:
|
||||
run:
|
||||
working-directory: nodejs
|
||||
@@ -169,19 +202,15 @@ jobs:
|
||||
# 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.
|
||||
# Docker builds can use rust-cache too: the workspace is bind-mounted, so
|
||||
# `target/` lives on the host and rust-cache's prune keeps the entry
|
||||
# small.
|
||||
#
|
||||
# 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.
|
||||
# bind-mounted from `.cargo-cache` rather than ~/.cargo, so that is cached
|
||||
# explicitly. And the key uses the *host* rustc version, not the compiler
|
||||
# that built these artifacts -- safe, since cargo fingerprints the real
|
||||
# one; a base-image bump just costs one cold build.
|
||||
- name: Cache cargo (docker builds)
|
||||
uses: Swatinem/rust-cache@v2
|
||||
if: ${{ matrix.settings.docker }}
|
||||
@@ -210,9 +239,14 @@ jobs:
|
||||
# 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.
|
||||
#
|
||||
# `docker run` inherits nothing; `-e NAME` carries the job's `env:` in.
|
||||
options: "--user 0:0 -v ${{ github.workspace }}/.cargo-cache/git/db:/usr/local/cargo/git/db \
|
||||
-v ${{ github.workspace }}/.cargo-cache/registry/cache:/usr/local/cargo/registry/cache \
|
||||
-v ${{ github.workspace }}/.cargo-cache/registry/index:/usr/local/cargo/registry/index \
|
||||
-e CARGO_PROFILE_RELEASE_LTO \
|
||||
-e CARGO_PROFILE_RELEASE_CODEGEN_UNITS \
|
||||
-e CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER \
|
||||
-v ${{ github.workspace }}:/build -w /build/nodejs"
|
||||
run: |
|
||||
set -e
|
||||
@@ -256,6 +290,18 @@ jobs:
|
||||
if: always()
|
||||
run: df -h
|
||||
shell: bash
|
||||
- name: Report peak memory
|
||||
if: always() && runner.os == 'Linux'
|
||||
shell: bash
|
||||
run: |
|
||||
peak=$(find /sys/fs/cgroup -name memory.peak -readable \
|
||||
-exec cat {} + 2>/dev/null | sort -n | tail -1)
|
||||
if [ -n "$peak" ]; then
|
||||
echo "peak memory: $((peak / 1024 / 1024)) MiB"
|
||||
else
|
||||
echo "peak memory: unavailable (no readable cgroup v2 memory.peak)"
|
||||
fi
|
||||
free -g || true
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
|
||||
Generated
+3
-3
@@ -5402,7 +5402,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.11"
|
||||
version = "0.38.0-beta.12"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"anyhow",
|
||||
@@ -5490,7 +5490,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-nodejs"
|
||||
version = "0.38.0-beta.11"
|
||||
version = "0.38.0-beta.12"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5515,7 +5515,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.11"
|
||||
version = "0.38.0-beta.12"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.38.0-beta.11</version>
|
||||
<version>0.38.0-beta.12</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.11</version>
|
||||
<version>0.38.0-beta.12</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+1
-1
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.11</version>
|
||||
<version>0.38.0-beta.12</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>${project.artifactId}</name>
|
||||
<description>LanceDB Java SDK Parent POM</description>
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.38.0-beta.11"
|
||||
version = "0.38.0-beta.12"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
+2
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
|
||||
+1
-1
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.38.0-beta.11",
|
||||
"version": "0.38.0-beta.12",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.11"
|
||||
version = "0.38.0-beta.12"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
|
||||
@@ -559,12 +559,18 @@ def _coerce_blob_list_values(
|
||||
|
||||
|
||||
def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
|
||||
if pa.types.is_null(values.type):
|
||||
data = pa.nulls(len(values), type=pa.large_binary())
|
||||
if _is_string_like(values.type):
|
||||
carrier_name = "uri"
|
||||
carrier = values
|
||||
elif pa.types.is_null(values.type):
|
||||
carrier_name = None
|
||||
carrier = None
|
||||
elif pa.types.is_large_binary(values.type):
|
||||
data = values
|
||||
carrier_name = "data"
|
||||
carrier = values
|
||||
else:
|
||||
data = values.cast(pa.large_binary())
|
||||
carrier_name = "data"
|
||||
carrier = values.cast(pa.large_binary())
|
||||
length = len(values)
|
||||
storage_type = target_field.type
|
||||
if isinstance(storage_type, pa.ExtensionType):
|
||||
@@ -572,8 +578,8 @@ def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
|
||||
storage_fields = list(storage_type)
|
||||
children = []
|
||||
for storage_field in storage_fields:
|
||||
if storage_field.name == "data":
|
||||
children.append(data)
|
||||
if storage_field.name == carrier_name:
|
||||
children.append(carrier.cast(storage_field.type))
|
||||
else:
|
||||
children.append(pa.nulls(length, type=storage_field.type))
|
||||
storage = pa.StructArray.from_arrays(
|
||||
@@ -593,7 +599,11 @@ def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]:
|
||||
|
||||
|
||||
def _can_coerce_to_blob(data_type: pa.DataType) -> bool:
|
||||
return _is_binary_like(data_type) or pa.types.is_null(data_type)
|
||||
return (
|
||||
_is_binary_like(data_type)
|
||||
or _is_string_like(data_type)
|
||||
or pa.types.is_null(data_type)
|
||||
)
|
||||
|
||||
|
||||
def _is_binary_like(data_type: pa.DataType) -> bool:
|
||||
@@ -604,6 +614,15 @@ def _is_binary_like(data_type: pa.DataType) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def _is_string_like(data_type: pa.DataType) -> bool:
|
||||
predicates = ("is_string", "is_large_string", "is_string_view")
|
||||
return any(
|
||||
predicate(data_type)
|
||||
for name in predicates
|
||||
if (predicate := getattr(pa.types, name, None)) is not None
|
||||
)
|
||||
|
||||
|
||||
def _field_extension_name(field: pa.Field) -> Optional[str]:
|
||||
extension_name = getattr(field.type, "extension_name", None)
|
||||
if extension_name is not None:
|
||||
@@ -618,6 +637,187 @@ def _field_extension_name(field: pa.Field) -> Optional[str]:
|
||||
return extension_name
|
||||
|
||||
|
||||
_JSON_EXTENSION_NAMES = {"arrow.json", "lance.json"}
|
||||
_BLOB_EXTENSION_NAME = "lance.blob.v2"
|
||||
|
||||
|
||||
def _field_contains_write_extension(field: pa.Field) -> bool:
|
||||
extension_name = _field_extension_name(field)
|
||||
if (
|
||||
extension_name in _JSON_EXTENSION_NAMES
|
||||
or extension_name == _BLOB_EXTENSION_NAME
|
||||
):
|
||||
return True
|
||||
if pa.types.is_struct(field.type):
|
||||
return any(_field_contains_write_extension(child) for child in field.type)
|
||||
if (
|
||||
pa.types.is_list(field.type)
|
||||
or pa.types.is_large_list(field.type)
|
||||
or pa.types.is_fixed_size_list(field.type)
|
||||
):
|
||||
return _field_contains_write_extension(field.type.value_field)
|
||||
return False
|
||||
|
||||
|
||||
def _with_field_type(
|
||||
field: pa.Field,
|
||||
data_type: pa.DataType,
|
||||
*,
|
||||
name: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> pa.Field:
|
||||
return pa.field(
|
||||
name or field.name,
|
||||
data_type,
|
||||
nullable=field.nullable,
|
||||
metadata=field.metadata if metadata is None else metadata,
|
||||
)
|
||||
|
||||
|
||||
def _with_list_value_field(
|
||||
data_type: pa.DataType, value_field: pa.Field
|
||||
) -> pa.DataType:
|
||||
if pa.types.is_list(data_type):
|
||||
return pa.list_(value_field)
|
||||
if pa.types.is_large_list(data_type):
|
||||
return pa.large_list(value_field)
|
||||
return pa.list_(value_field, data_type.list_size)
|
||||
|
||||
|
||||
def _extension_storage_field(field: pa.Field) -> pa.Field:
|
||||
"""Return a from-pylist-compatible field for nested write extensions."""
|
||||
extension_name = _field_extension_name(field)
|
||||
if extension_name in _JSON_EXTENSION_NAMES:
|
||||
metadata = dict(field.metadata or {})
|
||||
metadata[b"ARROW:extension:name"] = b"arrow.json"
|
||||
return _with_field_type(field, pa.string(), metadata=metadata)
|
||||
if extension_name == _BLOB_EXTENSION_NAME:
|
||||
metadata = dict(field.metadata or {})
|
||||
metadata[b"ARROW:extension:name"] = _BLOB_EXTENSION_NAME.encode()
|
||||
metadata[b"ARROW:extension:metadata"] = b""
|
||||
storage_type = getattr(field.type, "storage_type", field.type)
|
||||
return _with_field_type(field, storage_type, metadata=metadata)
|
||||
if pa.types.is_struct(field.type):
|
||||
children = [_extension_storage_field(child) for child in field.type]
|
||||
return _with_field_type(field, pa.struct(children))
|
||||
if _is_list_like(field.type):
|
||||
value_field = _extension_storage_field(field.type.value_field)
|
||||
return _with_field_type(field, _with_list_value_field(field.type, value_field))
|
||||
return field
|
||||
|
||||
|
||||
def _prepare_extension_field(
|
||||
field: pa.Field, target_field: pa.Field
|
||||
) -> Tuple[pa.Field, bool]:
|
||||
extension_name = _field_extension_name(target_field)
|
||||
if extension_name in _JSON_EXTENSION_NAMES:
|
||||
metadata = dict(field.metadata or {})
|
||||
metadata[b"ARROW:extension:name"] = b"arrow.json"
|
||||
return _with_field_type(field, pa.string(), metadata=metadata), True
|
||||
if extension_name == _BLOB_EXTENSION_NAME and pa.types.is_null(field.type):
|
||||
return _with_field_type(field, pa.large_binary()), True
|
||||
|
||||
if pa.types.is_struct(field.type) and pa.types.is_struct(target_field.type):
|
||||
target_children = {child.name: child for child in target_field.type}
|
||||
children = []
|
||||
changed = False
|
||||
for child in field.type:
|
||||
target_child = target_children.get(child.name)
|
||||
if target_child is None:
|
||||
children.append(child)
|
||||
continue
|
||||
prepared, child_changed = _prepare_extension_field(child, target_child)
|
||||
children.append(prepared)
|
||||
changed = changed or child_changed
|
||||
if changed:
|
||||
return _with_field_type(field, pa.struct(children)), True
|
||||
|
||||
if _is_list_like(field.type) and _is_list_like(target_field.type):
|
||||
target_value_field = target_field.type.value_field
|
||||
if _field_contains_write_extension(target_value_field):
|
||||
prepared = _extension_storage_field(target_value_field)
|
||||
data_type = _with_list_value_field(target_field.type, prepared)
|
||||
return _with_field_type(field, data_type), True
|
||||
|
||||
return field, False
|
||||
|
||||
|
||||
def _prepare_extension_value(
|
||||
value: Any, target_field: pa.Field, *, within_list: bool = False
|
||||
) -> Any:
|
||||
"""Shape raw nested blob values for PyArrow's struct construction."""
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
extension_name = _field_extension_name(target_field)
|
||||
if extension_name == _BLOB_EXTENSION_NAME and within_list:
|
||||
if isinstance(value, (bytes, bytearray, memoryview)):
|
||||
return {"data": value}
|
||||
if isinstance(value, str):
|
||||
return {"uri": value}
|
||||
return value
|
||||
|
||||
if pa.types.is_struct(target_field.type) and isinstance(value, dict):
|
||||
target_children = {child.name: child for child in target_field.type}
|
||||
return {
|
||||
name: _prepare_extension_value(
|
||||
child_value, target_children[name], within_list=within_list
|
||||
)
|
||||
if name in target_children
|
||||
else child_value
|
||||
for name, child_value in value.items()
|
||||
}
|
||||
|
||||
if _is_list_like(target_field.type) and isinstance(value, (list, tuple)):
|
||||
return [
|
||||
_prepare_extension_value(
|
||||
item, target_field.type.value_field, within_list=True
|
||||
)
|
||||
for item in value
|
||||
]
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def _prepare_extension_list(data: DATA, target_schema: pa.Schema) -> DATA:
|
||||
"""Give inferred list columns the logical type required by extensions."""
|
||||
if not isinstance(data, list) or not data or not isinstance(data[0], dict):
|
||||
return data
|
||||
|
||||
target_fields = {field.name: field for field in target_schema}
|
||||
if not any(
|
||||
_field_contains_write_extension(field) for field in target_fields.values()
|
||||
):
|
||||
return data
|
||||
|
||||
inferred = pa.Table.from_pylist(data)
|
||||
fields = []
|
||||
changed = False
|
||||
for field in inferred.schema:
|
||||
target_field = target_fields.get(field.name)
|
||||
if target_field is None:
|
||||
fields.append(field)
|
||||
continue
|
||||
prepared, field_changed = _prepare_extension_field(field, target_field)
|
||||
fields.append(prepared)
|
||||
changed = changed or field_changed
|
||||
|
||||
if not changed:
|
||||
return inferred
|
||||
|
||||
insert_schema = pa.schema(fields, metadata=inferred.schema.metadata)
|
||||
prepared_data = [
|
||||
{
|
||||
name: _prepare_extension_value(value, target_fields[name])
|
||||
if name in target_fields
|
||||
else value
|
||||
for name, value in row.items()
|
||||
}
|
||||
for row in data
|
||||
]
|
||||
return pa.Table.from_pylist(prepared_data, schema=insert_schema)
|
||||
|
||||
|
||||
def _align_field_types(
|
||||
fields: List[pa.Field],
|
||||
target_fields: List[pa.Field],
|
||||
@@ -5624,6 +5824,9 @@ class AsyncTable:
|
||||
if fill_value is None:
|
||||
fill_value = 0.0
|
||||
|
||||
if mode != "overwrite":
|
||||
data = _prepare_extension_list(data, schema)
|
||||
|
||||
# _santitize_data is an old code path, but we will use it until the
|
||||
# new code path is ready.
|
||||
if mode == "overwrite":
|
||||
|
||||
@@ -710,6 +710,80 @@ def test_fetch_blobs_preserves_null_and_empty_values():
|
||||
assert blobs[3].as_py() == b"present"
|
||||
|
||||
|
||||
def test_add_all_null_list_to_blob_column():
|
||||
table = _blob_table("all_null_add", [{"id": 1, "image": None}])
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert len(blobs) == 1
|
||||
assert blobs[0].as_py() is None
|
||||
|
||||
|
||||
def test_add_all_null_list_to_blob_column_with_sanitizer():
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("all_null_sanitized_add", schema=schema)
|
||||
|
||||
table.add([{"id": 1, "image": None}], on_bad_vectors="fill")
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert len(blobs) == 1
|
||||
assert blobs[0].as_py() is None
|
||||
|
||||
|
||||
def test_add_all_null_list_to_nested_blob_column():
|
||||
db = lancedb.connect("memory:///")
|
||||
blob_field = lancedb.blob("image")
|
||||
info_field = pa.field("info", pa.struct([blob_field]))
|
||||
info = pa.StructArray.from_arrays(
|
||||
[_blob_array("image", [b"seed"])], fields=[blob_field]
|
||||
)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), info],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
|
||||
)
|
||||
table = db.create_table("nested_null_add", data=seed)
|
||||
|
||||
table.add([{"id": 1, "info": {"image": None}}])
|
||||
table.add([{"id": 2, "info": {"image": None}}], on_bad_vectors="fill")
|
||||
|
||||
hits = table.search().where("id > 0").to_arrow()
|
||||
blobs = table.fetch_blobs("info.image", hits)
|
||||
assert len(blobs) == 2
|
||||
assert all(blob.as_py() is None for blob in blobs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("large_list", [False, True], ids=["list", "large_list"])
|
||||
def test_add_list_of_dicts_to_blob_list_column(large_list):
|
||||
db = lancedb.connect("memory:///")
|
||||
blob_field = lancedb.blob("image")
|
||||
blob_values = _blob_array("image", [b"seed"])
|
||||
if large_list:
|
||||
items_field = pa.field("items", pa.large_list(blob_field))
|
||||
items = pa.LargeListArray.from_arrays(
|
||||
pa.array([0, 1], type=pa.int64()), blob_values
|
||||
)
|
||||
else:
|
||||
items_field = pa.field("items", pa.list_(blob_field))
|
||||
items = pa.ListArray.from_arrays(pa.array([0, 1], type=pa.int32()), blob_values)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), items],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), items_field]),
|
||||
)
|
||||
table = db.create_table(f"blob_{large_list}_list_add", data=seed)
|
||||
|
||||
table.add([{"id": 1, "items": [None]}])
|
||||
table.add(
|
||||
[{"id": 2, "items": [b"a", None]}],
|
||||
on_bad_vectors="fill",
|
||||
)
|
||||
|
||||
ids = table.search().select(["id"]).to_arrow()["id"].to_pylist()
|
||||
assert sorted(ids) == [0, 1, 2]
|
||||
assert pa.types.is_large_list(table.schema.field("items").type) is large_list
|
||||
|
||||
|
||||
def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
|
||||
table = _blob_table(
|
||||
"range_alignment",
|
||||
@@ -1230,6 +1304,7 @@ def test_add_external_uri_string_round_trips_with_flag(tmp_path):
|
||||
table = db.create_table("external_string", schema=schema)
|
||||
table.add(
|
||||
[{"id": 1, "image": blob_path.as_uri()}],
|
||||
on_bad_vectors="fill",
|
||||
allow_external_blob_outside_bases=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -786,6 +786,55 @@ async def test_add_async(mem_db_async: AsyncConnection):
|
||||
assert await table.count_rows() == 3
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("values", "expected"),
|
||||
[
|
||||
([None], [None]),
|
||||
([None, '{"k": 1}'], [None, '{"k":1}']),
|
||||
(['{"k": 2}'], ['{"k":2}']),
|
||||
],
|
||||
)
|
||||
async def test_add_list_of_dicts_to_json_column(
|
||||
mem_db_async: AsyncConnection, values, expected
|
||||
):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), pa.field("value", pa.json_())])
|
||||
table = await mem_db_async.create_table("json_list_add", schema=schema)
|
||||
|
||||
await table.add([{"id": idx, "value": value} for idx, value in enumerate(values)])
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert [row["value"] for row in rows] == expected
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_list_of_dicts_to_nested_json_column(
|
||||
mem_db_async: AsyncConnection,
|
||||
):
|
||||
json_field = pa.field("value", pa.json_())
|
||||
info_field = pa.field("info", pa.struct([json_field]))
|
||||
info = pa.StructArray.from_arrays(
|
||||
[pa.array(['{"seed": 0}'], type=pa.json_())], fields=[json_field]
|
||||
)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), info],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
|
||||
)
|
||||
table = await mem_db_async.create_table("nested_json_list_add", data=seed)
|
||||
|
||||
await table.add([{"id": 1, "info": {"value": '{"k": 1}'}}])
|
||||
await table.add([{"id": 2, "info": {"value": '{"k": 2}'}}], on_bad_vectors="fill")
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert rows == [
|
||||
{"id": 0, "info": {"value": '{"seed":0}'}},
|
||||
{"id": 1, "info": {"value": '{"k":1}'}},
|
||||
{"id": 2, "info": {"value": '{"k":2}'}},
|
||||
]
|
||||
|
||||
|
||||
def test_add_overwrite_infers_vector_schema(mem_db: DBConnection):
|
||||
"""Overwrite should infer vector columns the same way create_table does.
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.11"
|
||||
version = "0.38.0-beta.12"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
|
||||
@@ -22,7 +22,7 @@
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef};
|
||||
use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef};
|
||||
use datafusion_common::tree_node::TreeNode;
|
||||
use datafusion_physical_plan::PhysicalExpr;
|
||||
use lance::dataset::NewColumnTransform;
|
||||
@@ -1273,6 +1273,11 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
|
||||
/// refresh time: that the expression parses, that every column it reads
|
||||
/// exists, and that the target name is free. A declaration that survives this
|
||||
/// is one a refresh can always act on.
|
||||
///
|
||||
/// Each accepted column joins the schema the next one resolves against, so a
|
||||
/// batch may declare `a` and then `b = a + 1` in one commit. Refresh order
|
||||
/// then matters, and refresh enforces it: `b` is refused while `a` still has
|
||||
/// unfilled rows.
|
||||
pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Vec<ArrowField>> {
|
||||
if columns.is_empty() {
|
||||
return Err(Error::InvalidInput {
|
||||
@@ -1280,11 +1285,11 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
|
||||
});
|
||||
}
|
||||
|
||||
let mut schema = schema;
|
||||
let mut fields = Vec::with_capacity(columns.len());
|
||||
let mut declared: Vec<&str> = Vec::with_capacity(columns.len());
|
||||
|
||||
for (name, expression) in columns {
|
||||
if schema.field_with_name(name).is_ok() || declared.contains(&name.as_str()) {
|
||||
if schema.field_with_name(name).is_ok() {
|
||||
return Err(Error::ColumnAlreadyExists { name: name.clone() });
|
||||
}
|
||||
|
||||
@@ -1292,16 +1297,50 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
|
||||
|
||||
// Declared columns start entirely null, so nullability is a property
|
||||
// of the declaration rather than of what the expression yields.
|
||||
fields.push(
|
||||
ArrowField::new(name, bound.data_type, true)
|
||||
.with_metadata(computed_column_metadata(expression, &bound.inputs)),
|
||||
);
|
||||
declared.push(name);
|
||||
let field = ArrowField::new(name, bound.data_type, true)
|
||||
.with_metadata(computed_column_metadata(expression, &bound.inputs));
|
||||
schema = Arc::new(ArrowSchema::new_with_metadata(
|
||||
schema
|
||||
.fields()
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain(std::iter::once(Arc::new(field.clone())))
|
||||
.collect::<Fields>(),
|
||||
schema.metadata().clone(),
|
||||
));
|
||||
fields.push(field);
|
||||
}
|
||||
|
||||
Ok(fields)
|
||||
}
|
||||
|
||||
/// Run the schema-level checks of
|
||||
/// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) against
|
||||
/// `schema` without committing: the Function-binding guard and the planning of
|
||||
/// every declaration. For callers that stage declarations behind other work
|
||||
/// and need those rejections before any of it lands.
|
||||
///
|
||||
/// Only the schema is consulted. Declaring also refuses a table with an LSM
|
||||
/// write spec or retained SSTables; that is table state, checked at commit.
|
||||
///
|
||||
/// ```
|
||||
/// # use std::sync::Arc;
|
||||
/// # use arrow_schema::{DataType, Field, Schema};
|
||||
/// use lancedb::table::computed_columns::validate_declarations;
|
||||
///
|
||||
/// let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
|
||||
/// let declarations = vec![
|
||||
/// ("a".to_string(), "x + 1".to_string()),
|
||||
/// ("b".to_string(), "a * 2".to_string()),
|
||||
/// ];
|
||||
/// assert!(validate_declarations(schema.clone(), &declarations).is_ok());
|
||||
/// assert!(validate_declarations(schema, &[("c".into(), "random()".into())]).is_err());
|
||||
/// ```
|
||||
pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<()> {
|
||||
ensure_no_function_bindings_for_mutation(schema.as_ref(), "schema evolution")?;
|
||||
plan(schema, columns).map(drop)
|
||||
}
|
||||
|
||||
/// Build the transform that declares `columns` against `schema`.
|
||||
///
|
||||
/// An all-null column is how a binding with no values yet is carried into a
|
||||
@@ -1340,6 +1379,22 @@ pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &st
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
/// The gate's reproducer: the validator applies the same schema-level
|
||||
/// guard declaring does, so a staging caller is refused before it commits
|
||||
/// anything else.
|
||||
#[test]
|
||||
fn test_validate_declarations_matches_schema_admission_barriers() {
|
||||
let schema = Arc::new(ArrowSchema::new_with_metadata(
|
||||
vec![ArrowField::new("x", DataType::Int32, true)],
|
||||
HashMap::from([(
|
||||
FUNCTION_BINDINGS_META_KEY.to_string(),
|
||||
"not valid binding metadata".to_string(),
|
||||
)]),
|
||||
));
|
||||
let declarations = vec![("a".to_string(), "x + 1".to_string())];
|
||||
assert!(super::validate_declarations(schema, &declarations).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_arrow_type_grammar_matches_the_shared_golden() {
|
||||
let golden: serde_json::Value = serde_json::from_str(include_str!(
|
||||
@@ -1582,6 +1637,40 @@ mod tests {
|
||||
assert!(declared(&table).await.is_empty());
|
||||
}
|
||||
|
||||
/// A batch may build on itself: one commit, and the later entry's inputs
|
||||
/// name the earlier one.
|
||||
#[tokio::test]
|
||||
async fn test_a_declaration_may_read_one_declared_before_it() {
|
||||
let table = table_with_ints("chain").await;
|
||||
let before = table.version().await.unwrap();
|
||||
add_computed(
|
||||
&table,
|
||||
&[("a".into(), "x + 1".into()), ("b".into(), "a * 2".into())],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(table.version().await.unwrap(), before + 1);
|
||||
let declared = declared(&table).await;
|
||||
assert_eq!(declared[1].name, "b");
|
||||
assert_eq!(declared[1].inputs, vec!["a".to_string()]);
|
||||
|
||||
// Order is the dependency order; reading ahead is still unknown.
|
||||
let err = add_computed(
|
||||
&table,
|
||||
&[("c".into(), "d + 1".into()), ("d".into(), "x + 1".into())],
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "c"));
|
||||
assert!(
|
||||
validate_declarations(
|
||||
table.schema().await.unwrap(),
|
||||
&[("e".into(), "random()".into())]
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
/// A column added by an ordinary transform is materialized, not bound, so
|
||||
/// it carries no declaration to report.
|
||||
#[tokio::test]
|
||||
|
||||
@@ -7,6 +7,16 @@
|
||||
//! therefore idempotent and does not observe input mutation -- once a row is
|
||||
//! filled, changing what the expression reads leaves the stored result alone.
|
||||
//!
|
||||
//! A column's computed inputs are filled first -- the dependency graph is
|
||||
//! walked once, each reachable column filled once in dependency order, each
|
||||
//! fill its own commit. Every fill in the pass, the requested column's
|
||||
//! included, covers only the fragments of the snapshot the pass started
|
||||
//! from: a commit may rebase over a concurrent append, and the fragment that
|
||||
//! admits carries placeholder nulls no earlier fill covered, so it waits for
|
||||
//! a later refresh rather than being read as values. Two concurrent fills of
|
||||
//! one input collide on its field in lance's conflict check, so a dependent
|
||||
//! fill can only commit over inputs that were durable when it read them.
|
||||
//!
|
||||
//! Two passes per fragment. The first scans only the unfilled live rows and
|
||||
//! evaluates the expression over them, which yields the exact fill count and
|
||||
//! decides whether the fragment is staged at all -- a fragment where nothing
|
||||
@@ -41,7 +51,8 @@ use crate::{Error, Result};
|
||||
/// The result of refreshing a computed column.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
pub struct RefreshColumnResult {
|
||||
/// Rows that had a value computed.
|
||||
/// Rows that had a value computed, in the requested column only; inputs
|
||||
/// filled on its behalf are not counted.
|
||||
#[serde(default)]
|
||||
pub rows_filled: u64,
|
||||
/// The commit version associated with the operation.
|
||||
@@ -52,6 +63,7 @@ pub struct RefreshColumnResult {
|
||||
struct RefreshExecution {
|
||||
result: RefreshColumnResult,
|
||||
source_version: u64,
|
||||
published_version: Option<u64>,
|
||||
}
|
||||
|
||||
/// Internal implementation of the refresh logic.
|
||||
@@ -74,7 +86,12 @@ async fn execute_refresh_column_with_source(
|
||||
|
||||
let expression = declared_expression(&dataset, column)?;
|
||||
let schema = Arc::new(ArrowSchema::from(dataset.schema()));
|
||||
let bound = Arc::new(super::computed_columns::bind(schema, column, &expression)?);
|
||||
let bound = Arc::new(super::computed_columns::bind(
|
||||
schema.clone(),
|
||||
column,
|
||||
&expression,
|
||||
)?);
|
||||
ensure_inputs_filled(&dataset, &schema, column, &bound).await?;
|
||||
let field = dataset
|
||||
.schema()
|
||||
.field(column)
|
||||
@@ -100,25 +117,25 @@ async fn execute_refresh_column_with_source(
|
||||
replacements.push(fragment.write_columns(values, &column_schema).await?);
|
||||
}
|
||||
|
||||
let source_version = dataset.version().version;
|
||||
if replacements.is_empty() {
|
||||
let source_version = dataset.version().version;
|
||||
return Ok(RefreshExecution {
|
||||
result: RefreshColumnResult {
|
||||
rows_filled: 0,
|
||||
version: source_version,
|
||||
},
|
||||
source_version,
|
||||
published_version: None,
|
||||
});
|
||||
}
|
||||
|
||||
let read_version = dataset.version().version;
|
||||
// The dataset's own session, so registrations and caches survive the
|
||||
// commit being installed on the handle.
|
||||
let session = dataset.session();
|
||||
let new_dataset = Dataset::commit(
|
||||
WriteDestination::Dataset(dataset.clone()),
|
||||
Operation::DataReplacement { replacements },
|
||||
Some(read_version),
|
||||
Some(source_version),
|
||||
None,
|
||||
None,
|
||||
session,
|
||||
@@ -133,10 +150,52 @@ async fn execute_refresh_column_with_source(
|
||||
rows_filled,
|
||||
version,
|
||||
},
|
||||
source_version: read_version,
|
||||
source_version,
|
||||
published_version: Some(version),
|
||||
})
|
||||
}
|
||||
|
||||
/// Refuse while a computed input still has rows a refresh of it would fill:
|
||||
/// read now, its placeholder null would be evaluated as a value and kept.
|
||||
async fn ensure_inputs_filled(
|
||||
dataset: &Dataset,
|
||||
schema: &Arc<ArrowSchema>,
|
||||
column: &str,
|
||||
bound: &BoundExpression,
|
||||
) -> Result<()> {
|
||||
for input in &bound.roots {
|
||||
let Some(declaration) = schema
|
||||
.field_with_name(input)
|
||||
.ok()
|
||||
.and_then(computed_column_from_field)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let ComputedColumnKind::Sql { expression } = &declaration.kind else {
|
||||
return Err(Error::NotSupported {
|
||||
message: format!(
|
||||
"computed column '{column}' reads '{input}', whose fill state this \
|
||||
refresh cannot check; refresh '{input}' first"
|
||||
),
|
||||
});
|
||||
};
|
||||
let input_bound = super::computed_columns::bind(schema.clone(), input, expression)?;
|
||||
let mut unfilled = 0u64;
|
||||
for fragment in dataset.get_fragments() {
|
||||
unfilled += count_fragment_gains(dataset, &fragment, &input_bound, input).await?;
|
||||
}
|
||||
if unfilled > 0 {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"computed column '{column}' reads '{input}', which has {unfilled} unfilled \
|
||||
rows; refresh '{input}' first"
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run the refresh as a [`Job`] in this process.
|
||||
pub(crate) async fn execute_refresh_column_async(
|
||||
table: &NativeTable,
|
||||
@@ -160,8 +219,7 @@ pub(crate) async fn execute_refresh_column_async(
|
||||
rows_failed: 0,
|
||||
rows_remaining: 0,
|
||||
source_version: execution.source_version,
|
||||
published_version: (execution.result.rows_filled > 0)
|
||||
.then_some(execution.result.version),
|
||||
published_version: execution.published_version,
|
||||
})
|
||||
})))
|
||||
}
|
||||
@@ -384,7 +442,8 @@ mod tests {
|
||||
.version)
|
||||
}
|
||||
|
||||
async fn read(table: &Table, column: &str) -> Vec<Option<i32>> {
|
||||
async fn read(table: &Table, column: &str) -> Vec<Option<i64>> {
|
||||
use arrow_array::{Array, Int64Array};
|
||||
let batches = table
|
||||
.query()
|
||||
.select(Select::columns(&[column]))
|
||||
@@ -394,15 +453,19 @@ mod tests {
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let mut values: Vec<Option<i32>> = batches
|
||||
let mut values: Vec<Option<i64>> = batches
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch[column]
|
||||
.as_any()
|
||||
.downcast_ref::<Int32Array>()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.collect::<Vec<_>>()
|
||||
let array = &batch[column];
|
||||
match array.as_any().downcast_ref::<Int32Array>() {
|
||||
Some(ints) => ints.iter().map(|v| v.map(i64::from)).collect::<Vec<_>>(),
|
||||
None => array
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.collect::<Vec<_>>(),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
values.sort();
|
||||
@@ -414,6 +477,98 @@ mod tests {
|
||||
table.add(batch).execute().await.unwrap();
|
||||
}
|
||||
|
||||
/// The gate's reproducer: `b = coalesce(a, 0)` refreshed before `a`
|
||||
/// must not bake zeros from `a`'s placeholder null. It is refused, and
|
||||
/// names the input, until `a` is filled -- after every append too.
|
||||
#[tokio::test]
|
||||
async fn test_dependent_refresh_refuses_an_unfilled_input() {
|
||||
let table = table_with("dependent_refresh_order", vec![1, 2, 3]).await;
|
||||
table
|
||||
.add_columns()
|
||||
.computed("a", "x + 1")
|
||||
.computed("b", "coalesce(a, 0)")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let err = table.refresh_column("b").await.unwrap_err();
|
||||
assert!(
|
||||
matches!(&err, Error::InvalidInput { message } if message.contains("refresh 'a' first")),
|
||||
"{err}"
|
||||
);
|
||||
assert_eq!(read(&table, "b").await, vec![None, None, None]);
|
||||
|
||||
assert_eq!(table.refresh_column("a").await.unwrap().rows_filled, 3);
|
||||
assert_eq!(table.refresh_column("b").await.unwrap().rows_filled, 3);
|
||||
assert_eq!(read(&table, "b").await, vec![Some(2), Some(3), Some(4)]);
|
||||
|
||||
append(&table, vec![10]).await;
|
||||
assert!(table.refresh_column("b").await.is_err());
|
||||
table.refresh_column("a").await.unwrap();
|
||||
assert_eq!(table.refresh_column("b").await.unwrap().rows_filled, 1);
|
||||
assert_eq!(
|
||||
table.count_rows(Some("b = 0".to_string())).await.unwrap(),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
/// Names that need quoting, and a nested input, survive the trip through
|
||||
/// declaration metadata and the dependency check: the recorded inputs
|
||||
/// are matched by name, never re-parsed as SQL.
|
||||
#[tokio::test]
|
||||
async fn test_dependent_refresh_handles_awkward_column_names() {
|
||||
use arrow_array::{Int32Array, StructArray};
|
||||
use arrow_schema::{DataType, Field, Fields};
|
||||
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let age_fields = Fields::from(vec![Field::new("age", DataType::Int32, true)]);
|
||||
let meta = StructArray::new(
|
||||
age_fields.clone(),
|
||||
vec![Arc::new(Int32Array::from(vec![10, 20])) as _],
|
||||
None,
|
||||
);
|
||||
let schema = Arc::new(arrow_schema::Schema::new(vec![
|
||||
Field::new("camelCase", DataType::Int32, true),
|
||||
Field::new("with-hyphen", DataType::Int32, true),
|
||||
Field::new("meta", DataType::Struct(age_fields), true),
|
||||
]));
|
||||
let batch = arrow_array::RecordBatch::try_new(
|
||||
schema,
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1, 2])) as _,
|
||||
Arc::new(Int32Array::from(vec![100, 200])) as _,
|
||||
Arc::new(meta) as _,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let table = conn
|
||||
.create_table("awkward_names", batch)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
table
|
||||
.add_columns()
|
||||
.computed("y", "`camelCase` * 2")
|
||||
.computed("z", "coalesce(y, 0) + `with-hyphen` + meta.age")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let z = crate::table::computed_columns::computed_columns(
|
||||
table.schema().await.unwrap().as_ref(),
|
||||
)
|
||||
.into_iter()
|
||||
.find(|c| c.name == "z")
|
||||
.unwrap();
|
||||
assert_eq!(z.inputs, vec!["meta.age", "with-hyphen", "y"]);
|
||||
|
||||
let err = table.refresh_column("z").await.unwrap_err();
|
||||
assert!(err.to_string().contains("refresh 'y' first"), "{err}");
|
||||
assert_eq!(table.refresh_column("y").await.unwrap().rows_filled, 2);
|
||||
assert_eq!(table.refresh_column("z").await.unwrap().rows_filled, 2);
|
||||
assert_eq!(read(&table, "z").await, vec![Some(112), Some(224)]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_fills_a_declared_column() {
|
||||
let table = table_with("refresh_fills", vec![1, 2, 3]).await;
|
||||
@@ -651,7 +806,8 @@ mod tests {
|
||||
|
||||
let read_back = read(&table, "doubled").await;
|
||||
assert_eq!(read_back.len(), 20_000);
|
||||
let mut expected: Vec<Option<i32>> = values.iter().map(|v| Some(v * 2)).collect();
|
||||
let mut expected: Vec<Option<i64>> =
|
||||
values.iter().map(|v| Some(i64::from(v * 2))).collect();
|
||||
expected.sort();
|
||||
assert_eq!(read_back, expected);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user