diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..fc38587579e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -162,7 +162,13 @@ flatbuffers = "25.2.10" fsst-rs = "0.6.0" futures = { version = "0.3.31", default-features = false } fuzzy-matcher = "0.3" -geo = "0.31.0" +# `vortex-spatial`'s `contains_route` transcribes geo's `impl_contains_from_relate!` dispatch +# table, so any bump that moves a row silently changes containment verdicts. The tests stay green +# wherever relate and the direct algorithm agree. Pinned exactly so that taking any new geo, +# patch releases included, is a deliberate edit of this line that re-verifies the table; a caret +# requirement would let `cargo update` (or automated lockfile maintenance) take 0.31.x with no diff +# to review. See `vortex-spatial/src/scalar_fn/contains.rs`. +geo = "=0.31.0" geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs index 96f2afff50a..be8fd0ba762 100644 --- a/encodings/runend/src/trace_tests.rs +++ b/encodings/runend/src/trace_tests.rs @@ -73,6 +73,14 @@ fn trace_compare_on_runend() -> VortexResult<()> { iter 0 current=vortex.runend(bool, len=9) builder_active=false execute_until target=AnyCanonical root=vortex.binary(bool, len=3) iter 0 current=vortex.binary(bool, len=3) builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 1 current=vortex.bool(bool, len=3) builder_active=false return output=vortex.bool(bool, len=3) diff --git a/scripts/benchmark-rowfn.sh b/scripts/benchmark-rowfn.sh new file mode 100755 index 00000000000..7536e37ec87 --- /dev/null +++ b/scripts/benchmark-rowfn.sh @@ -0,0 +1,488 @@ +#!/usr/bin/env bash + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +set -Eeu -o pipefail + +script_directory=$(dirname "$(realpath "${BASH_SOURCE[0]}")") + +usage() { + cat >&2 <<'EOF' +Usage: benchmark-rowfn.sh [OPTIONS] + +Options: + --suite NAME Select a preset or benchmark label. Repeatable; defaults to full. + --filter PATTERN Pass a Divan benchmark filter. Repeatable. + --build-only Build and record benchmark executables without measuring. + --measure-only Measure previously recorded benchmark executables without building. + --config NAME repository (16 CGUs/no LTO, default) or primary (1 CGU/fat LTO). + --target-root PATH Parent for reusable baseline and candidate Cargo targets. + --baseline-target PATH Reusable Cargo target for the baseline revision. + --candidate-target PATH + Reusable Cargo target for the candidate revision. + --codegen-units N Override the selected configuration. + --lto VALUE Override LTO with false, thin, or fat. + --rustflags FLAGS Override RUSTFLAGS; defaults to -C target-cpu=native. + --build-jobs N Jobs per concurrent revision build; defaults to 8 and cannot exceed 8. + --bench-cpu N Logical CPU used for every timed process; defaults to 4. + --warm-runs N Warm runs per revision; defaults to 2. + --measured-pairs N Alternating measured pairs; defaults to 7. + --sample-count N Divan sample count; defaults to 100. + --min-time SECONDS Divan minimum time; defaults to 0.25. + --max-time SECONDS Divan maximum time; defaults to 0.5. + --lock-file PATH Global timed-run lock; defaults to /tmp/vortex-rowfn-benchmark.lock. + --list-suites Print presets and benchmark labels, then exit. +EOF +} + +suite_catalog=( + "array-binary_ops|vortex-array|binary_ops|array,numeric,design-a-matrix,full" + "array-compare|vortex-array|compare|array,compare,full" + "array-row_fn_executor|vortex-array|row_fn_executor|array,framework,full" + "array-strict_validity|vortex-array|strict_validity|array,framework,full" + "array-like|vortex-array|like|array,full" + "array-take_filter|vortex-array|take_filter|array,full" + "array-varbinview_compact|vortex-array|varbinview_compact|array,full" + "tensor-l2_norm|vortex-tensor|l2_norm|tensor,full" + "tensor-inner_product|vortex-tensor|inner_product|tensor,full" + "tensor-cosine_similarity|vortex-tensor|cosine_similarity|tensor,full" + "tensor-normalized|vortex-tensor|normalized|tensor,full" + "spatial-binary_predicates|vortex-spatial|binary_predicates|spatial,full" + "spatial-distance|vortex-spatial|distance|spatial,full" + "spatial-envelope|vortex-spatial|envelope|spatial,full" + "spatial-predicate_bbox|vortex-spatial|predicate_bbox|spatial,full" +) + +requested_suites=() +filters=() +run_build=true +run_measure=true +configuration=repository +target_root= +baseline_target_override= +candidate_target_override= +codegen_units_override= +lto_override= +rustflags_override= +build_jobs=8 +bench_cpu=4 +warm_runs=2 +measured_pairs=7 +sample_count=100 +min_time=0.25 +max_time=0.5 +lock_file=/tmp/vortex-rowfn-benchmark.lock + +while [[ $# -gt 0 ]]; do + case $1 in + --suite) requested_suites+=("$2"); shift 2 ;; + --filter) filters+=("$2"); shift 2 ;; + --build-only) + if [[ $run_build == false ]]; then + echo "--build-only and --measure-only are mutually exclusive." >&2 + exit 1 + fi + run_measure=false + shift + ;; + --measure-only) + if [[ $run_measure == false ]]; then + echo "--build-only and --measure-only are mutually exclusive." >&2 + exit 1 + fi + run_build=false + shift + ;; + --config) configuration=$2; shift 2 ;; + --target-root) target_root=$2; shift 2 ;; + --baseline-target) baseline_target_override=$2; shift 2 ;; + --candidate-target) candidate_target_override=$2; shift 2 ;; + --codegen-units) codegen_units_override=$2; shift 2 ;; + --lto) lto_override=$2; shift 2 ;; + --rustflags) rustflags_override=$2; shift 2 ;; + --build-jobs) build_jobs=$2; shift 2 ;; + --bench-cpu) bench_cpu=$2; shift 2 ;; + --warm-runs) warm_runs=$2; shift 2 ;; + --measured-pairs) measured_pairs=$2; shift 2 ;; + --sample-count) sample_count=$2; shift 2 ;; + --min-time) min_time=$2; shift 2 ;; + --max-time) max_time=$2; shift 2 ;; + --lock-file) lock_file=$2; shift 2 ;; + --list-suites) + echo "Presets: full array framework numeric design-a-matrix compare tensor spatial" + printf '%s\n' "${suite_catalog[@]}" | cut -d '|' -f 1 + exit 0 + ;; + -h|--help) usage; exit 0 ;; + --*) echo "Unknown option: $1" >&2; usage; exit 1 ;; + *) break ;; + esac +done + +if [[ $# -ne 3 ]]; then + usage + exit 1 +fi +if [[ $(uname -m) != x86_64 ]]; then + echo "RowFn native performance decisions require an x86_64 host." >&2 + exit 1 +fi +if ((build_jobs < 1 || build_jobs > 8)); then + echo "--build-jobs must be between 1 and 8 so two builds cannot exceed 16 jobs." >&2 + exit 1 +fi +if [[ $run_measure == true ]]; then + command -v flock >/dev/null || { echo "benchmark-rowfn.sh requires flock." >&2; exit 1; } +fi + +baseline=$(realpath "$1") +candidate=$(realpath "$2") +output=$(realpath -m "$3") +if [[ -e $output ]]; then + echo "Output path already exists: $output" >&2 + exit 1 +fi + +case $configuration in + primary) codegen_units=1; lto=fat ;; + repository) codegen_units=16; lto=false ;; + *) echo "Unknown configuration: $configuration" >&2; exit 1 ;; +esac +codegen_units=${codegen_units_override:-$codegen_units} +lto=${lto_override:-$lto} +rustflags=${rustflags_override:--C target-cpu=native} + +if ((${#requested_suites[@]} == 0)); then + requested_suites=(full) +fi +selected_suites=() +declare -A selected_labels=() +for request in "${requested_suites[@]}"; do + matched=false + for entry in "${suite_catalog[@]}"; do + IFS='|' read -r label _ _ groups <<<"$entry" + if [[ $request == "$label" || ,$groups, == *,$request,* ]]; then + matched=true + if [[ -z ${selected_labels[$label]:-} ]]; then + selected_suites+=("$entry") + selected_labels[$label]=1 + fi + fi + done + if [[ $matched == false ]]; then + echo "Unknown suite or benchmark label: $request" >&2 + exit 1 + fi +done + +common_suites=() +skipped_suites=() +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label package bench _ <<<"$entry" + baseline_source="$baseline/$package/benches/$bench.rs" + candidate_source="$candidate/$package/benches/$bench.rs" + + if [[ -f $baseline_source && -f $candidate_source ]]; then + common_suites+=("$entry") + elif [[ -f $baseline_source ]]; then + skipped_suites+=("$label (baseline only)") + elif [[ -f $candidate_source ]]; then + skipped_suites+=("$label (candidate only)") + else + skipped_suites+=("$label (missing from both revisions)") + fi +done +if ((${#common_suites[@]} == 0)); then + echo "No requested benchmark targets exist in both revisions; no comparison is possible." >&2 + printf 'Skipped: %s\n' "${skipped_suites[@]}" >&2 + exit 1 +fi +selected_suites=("${common_suites[@]}") +if ((${#skipped_suites[@]} != 0)); then + printf 'Skipping one-sided benchmark target: %s\n' "${skipped_suites[@]}" >&2 +fi + +common_git_dir=$(git -C "$candidate" rev-parse --path-format=absolute --git-common-dir) +repository_root=$(dirname "$common_git_dir") +if [[ -n $target_root && (-n $baseline_target_override || -n $candidate_target_override) ]]; then + echo "--target-root cannot be combined with revision-specific target paths." >&2 + exit 1 +fi +if [[ -z $target_root ]]; then + target_root="$repository_root/target/rowfn-benchmark/$(basename "$output")" +fi +target_root=$(realpath -m "$target_root") +baseline_target=$(realpath -m "${baseline_target_override:-$target_root/baseline}") +candidate_target=$(realpath -m "${candidate_target_override:-$target_root/candidate}") +if [[ $baseline_target == "$candidate_target" ]]; then + echo "Baseline and candidate must use different Cargo target directories." >&2 + exit 1 +fi + +mkdir -p "$output" +if [[ $run_build == true ]]; then + mkdir -p "$output/build" "$baseline_target" "$candidate_target" +fi +if [[ $run_measure == true ]]; then + mkdir -p "$output/warm" "$output/measured" +fi +parser="$script_directory/rowfn_benchmark.py" + +{ + echo "RowFn benchmark machine record" + echo "Date: $(date --iso-8601=seconds)" + echo "Host: $(hostname)" + echo "Kernel: $(uname -srvmo)" + echo "Benchmark CPU: $bench_cpu" + echo "Configuration: $configuration" + echo "Cargo profile: bench, $codegen_units codegen units, LTO $lto" + echo "RUSTFLAGS: $rustflags" + echo "Warm runs: $warm_runs" + echo "Measured pairs: $measured_pairs" + echo "Divan: TSC timer, $sample_count samples, min $min_time s, max $max_time s" + if ((${#skipped_suites[@]} == 0)); then + echo "Skipped one-sided benchmark targets: none" + else + printf 'Skipped one-sided benchmark target: %s\n' "${skipped_suites[@]}" + fi + echo + echo "Baseline toolchain:" + (cd "$baseline" && rustc -vV && cargo -V) + echo + echo "Candidate toolchain:" + (cd "$candidate" && rustc -vV && cargo -V) + echo + lscpu + echo + rg -m1 '^microcode' /proc/cpuinfo || true + for path in \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/scaling_governor \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/energy_performance_preference \ + /sys/devices/system/cpu/cpufreq/boost; do + [[ -r $path ]] && echo "$path: $(<"$path")" + done +} >"$output/machine.txt" + +build_revision() { + local worktree=$1 + local target=$2 + local log=$3 + + ( + cd "$worktree" + export CARGO_TARGET_DIR=$target + export CARGO_PROFILE_BENCH_CODEGEN_UNITS=$codegen_units + export CARGO_PROFILE_BENCH_LTO=$lto + export RUSTFLAGS=$rustflags + for package in vortex-array vortex-tensor vortex-spatial; do + local command=(cargo bench --no-run -j "$build_jobs" -p "$package") + local has_bench=false + for entry in "${selected_suites[@]}"; do + IFS='|' read -r _ suite_package bench _ <<<"$entry" + if [[ $suite_package == "$package" ]]; then + command+=(--bench "$bench") + has_bench=true + fi + done + if [[ $has_bench == true ]]; then + "${command[@]}" + fi + done + ) >"$log" 2>&1 +} + +if [[ $run_build == true ]]; then + echo "Building baseline and candidate with $build_jobs jobs each." + build_revision "$baseline" "$baseline_target" "$output/build/baseline.txt" & + baseline_pid=$! + build_revision "$candidate" "$candidate_target" "$output/build/candidate.txt" & + candidate_pid=$! + baseline_status=0 + candidate_status=0 + wait "$baseline_pid" || baseline_status=$? + wait "$candidate_pid" || candidate_status=$? + if [[ $baseline_status -ne 0 || $candidate_status -ne 0 ]]; then + echo "Benchmark build failed; see $output/build/." >&2 + exit 1 + fi +fi + +find_benchmark() { + local target=$1 + local name=$2 + local binary + + binary=$(find "$target/release/deps" -maxdepth 1 -type f -executable -name "$name-*" \ + -printf '%T@ %p\n' | sort -nr | head -n 1 | cut -d ' ' -f 2-) + [[ -n $binary ]] || { echo "Cannot find benchmark $name under $target." >&2; exit 1; } + echo "$binary" +} + +declare -A baseline_binaries=() +declare -A candidate_binaries=() +build_settings=( + --setting "configuration=$configuration" + --setting "codegen_units=$codegen_units" + --setting "lto=$lto" + --setting "rustflags=$rustflags" +) + +record_build() { + local revision=$1 + local worktree=$2 + local target=$3 + local metadata="$target/rowfn-benchmark-build.json" + local arguments=( + record-build + --output "$metadata" + --worktree "$worktree" + --target "$target" + "${build_settings[@]}" + ) + + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ bench _ <<<"$entry" + local binary + binary=$(find_benchmark "$target" "$bench") + arguments+=(--binary "$label=$binary") + done + python3 "$parser" "${arguments[@]}" + echo "Recorded $revision build metadata: $metadata" +} + +if [[ $run_build == true ]]; then + record_build baseline "$baseline" "$baseline_target" + record_build candidate "$candidate" "$candidate_target" +fi + +load_binaries() { + local worktree=$1 + local target=$2 + local metadata="$target/rowfn-benchmark-build.json" + local arguments=( + validate-build + --metadata "$metadata" + --worktree "$worktree" + --target "$target" + "${build_settings[@]}" + ) + + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + arguments+=(--suite "$label") + done + python3 "$parser" "${arguments[@]}" +} + +baseline_binary_output=$(load_binaries "$baseline" "$baseline_target") +candidate_binary_output=$(load_binaries "$candidate" "$candidate_target") +mapfile -t baseline_binary_records <<<"$baseline_binary_output" +mapfile -t candidate_binary_records <<<"$candidate_binary_output" +for record in "${baseline_binary_records[@]}"; do + label=${record%%=*} + baseline_binaries[$label]=${record#*=} +done +for record in "${candidate_binary_records[@]}"; do + label=${record%%=*} + candidate_binaries[$label]=${record#*=} +done + +manifest_args=( + manifest + --output "$output/manifest.json" + --machine-record "$output/machine.txt" + --baseline-worktree "$baseline" + --candidate-worktree "$candidate" + --baseline-target "$baseline_target" + --candidate-target "$candidate_target" + --setting "configuration=$configuration" + --setting "codegen_units=$codegen_units" + --setting "lto=$lto" + --setting "rustflags=$rustflags" + --setting "bench_cpu=$bench_cpu" + --setting "warm_runs=$warm_runs" + --setting "measured_pairs=$measured_pairs" + --setting "sample_count=$sample_count" + --setting "min_time=$min_time" + --setting "max_time=$max_time" +) +for filter in "${filters[@]}"; do + manifest_args+=(--filter "$filter") +done +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + manifest_args+=( + --suite "$label" + --baseline-binary "$label=${baseline_binaries[$label]}" + --candidate-binary "$label=${candidate_binaries[$label]}" + ) +done +python3 "$parser" "${manifest_args[@]}" + +if [[ $run_measure == false ]]; then + echo "Build evidence: $output" + echo "Baseline target: $baseline_target" + echo "Candidate target: $candidate_target" + exit 0 +fi + +run_suite() { + local revision=$1 + local label=$2 + local destination=$3 + local binary + local command + + if [[ $revision == baseline ]]; then + binary=${baseline_binaries[$label]} + else + binary=${candidate_binaries[$label]} + fi + command=( + taskset -c "$bench_cpu" "$binary" + --bench --timer tsc --sample-count "$sample_count" + --min-time "$min_time" --max-time "$max_time" --color never + "${filters[@]}" + ) + echo "Running $label ($revision) -> $destination" + "${command[@]}" >"$destination" 2>&1 +} + +echo "Waiting for the global timed benchmark lock: $lock_file" +exec {benchmark_lock}>"$lock_file" +flock "$benchmark_lock" +if pgrep -x cargo >/dev/null || pgrep -x rustc >/dev/null; then + echo "Cargo or rustc is active after acquiring the benchmark lock; refusing to measure." >&2 + exit 1 +fi + +for ((round = 1; round <= warm_runs; round++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((round % 2 == 1)); then + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + else + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + fi + done +done + +for ((pair = 1; pair <= measured_pairs; pair++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((pair % 2 == 1)); then + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + else + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + fi + done +done + +python3 "$parser" summarize "$output" +echo "Raw results: $output" +echo "Summary: $output/summary.md" diff --git a/scripts/rowfn_benchmark.py b/scripts/rowfn_benchmark.py new file mode 100755 index 00000000000..35c08918d75 --- /dev/null +++ b/scripts/rowfn_benchmark.py @@ -0,0 +1,482 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Capture and summarize evidence from ``benchmark-rowfn.sh`` runs.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import re +import statistics +import subprocess +from collections.abc import Iterable +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path + +RESULT_FILE = re.compile(r"^(?P.+)-(?Pbaseline|candidate)-(?P\d+)\.txt$") +TREE_ROW = re.compile(r"^(?P(?:│ | )*)(?:├─ |╰─ )(?P.*)$") +TIMING = re.compile(r"(?P\d+(?:\.\d+)?)\s*(?Pps|ns|µs|us|ms|s)\s*$") +UNIT_TO_NS = { + "ps": 0.001, + "ns": 1.0, + "µs": 1_000.0, + "us": 1_000.0, + "ms": 1_000_000.0, + "s": 1_000_000_000.0, +} + + +@dataclass(frozen=True) +class BenchmarkSummary: + suite: str + benchmark: str + pairs: int + baseline_median_ns: float + candidate_median_ns: float + median_ratio: float + minimum_ratio: float + maximum_ratio: float + ratio_mad: float + + +def run_git(worktree: Path, *args: str, binary: bool = False) -> str | bytes: + """Run one read-only Git command in ``worktree``.""" + + result = subprocess.run( + ["git", "-C", str(worktree), *args], + check=True, + capture_output=True, + text=not binary, + ) + return result.stdout if binary else result.stdout.strip() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + + return digest.hexdigest() + + +def toolchain_record(worktree: Path) -> dict[str, str]: + """Capture the tools selected from a revision's working directory.""" + + def version(*command: str) -> str: + result = subprocess.run( + command, + cwd=worktree, + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + return {"rustc": version("rustc", "-vV"), "cargo": version("cargo", "-V")} + + +def revision_record(worktree: Path, target: Path, binaries: Iterable[str]) -> dict[str, object]: + """Describe the exact revision, dirty patch, targets, and benchmark executables.""" + + status = str(run_git(worktree, "status", "--short")).splitlines() + diff = run_git(worktree, "diff", "--binary", "HEAD", binary=True) + assert isinstance(diff, bytes) + + untracked = run_git(worktree, "ls-files", "--others", "--exclude-standard", "-z", binary=True) + assert isinstance(untracked, bytes) + dirty_digest = hashlib.sha256(diff) + dirty_digest.update(untracked) + for relative_path in filter(None, untracked.decode().split("\0")): + path = worktree / relative_path + if path.is_file(): + dirty_digest.update(relative_path.encode()) + dirty_digest.update(bytes.fromhex(sha256_file(path))) + + executable_records: dict[str, object] = {} + for entry in binaries: + label, separator, raw_path = entry.partition("=") + if not separator: + raise ValueError(f"expected LABEL=PATH for benchmark binary, got {entry!r}") + path = Path(raw_path).resolve() + executable_records[label] = { + "path": str(path), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + + return { + "worktree": str(worktree.resolve()), + "head": run_git(worktree, "rev-parse", "HEAD"), + "changed_paths": status, + "tracked_diff_sha256": hashlib.sha256(diff).hexdigest(), + "dirty_state_sha256": dirty_digest.hexdigest(), + "target": str(target.resolve()), + "binaries": executable_records, + } + + +def build_identity(worktree: Path, target: Path, settings: dict[str, str]) -> dict[str, object]: + revision = revision_record(worktree, target, []) + revision.pop("binaries") + return { + "settings": settings, + "toolchain": toolchain_record(worktree), + "revision": revision, + } + + +def binary_records(entries: Iterable[str]) -> dict[str, object]: + records: dict[str, object] = {} + for entry in entries: + label, separator, raw_path = entry.partition("=") + if not separator: + raise ValueError(f"expected LABEL=PATH for benchmark binary, got {entry!r}") + path = Path(raw_path).resolve() + records[label] = { + "path": str(path), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + return records + + +def write_build_record(args: argparse.Namespace) -> None: + output = Path(args.output) + worktree = Path(args.worktree) + target = Path(args.target) + settings = dict(setting.split("=", 1) for setting in args.setting) + identity = build_identity(worktree, target, settings) + binaries = binary_records(args.binary) + + if output.exists(): + previous = json.loads(output.read_text(encoding="utf-8")) + previous_identity = {key: previous.get(key) for key in identity} + if previous_identity == identity: + binaries = {**previous.get("binaries", {}), **binaries} + + record = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + **identity, + "binaries": binaries, + } + output.write_text(f"{json.dumps(record, indent=2, sort_keys=True)}\n", encoding="utf-8") + + +def validated_build_binaries(args: argparse.Namespace) -> dict[str, str]: + metadata = Path(args.metadata) + if not metadata.is_file(): + raise ValueError(f"build metadata does not exist: {metadata}") + + record = json.loads(metadata.read_text(encoding="utf-8")) + if record.get("schema_version") != 1: + raise ValueError(f"unsupported build metadata schema in {metadata}") + + settings = dict(setting.split("=", 1) for setting in args.setting) + identity = build_identity(Path(args.worktree), Path(args.target), settings) + mismatches = [key for key in identity if record.get(key) != identity[key]] + if mismatches: + fields = ", ".join(mismatches) + raise ValueError(f"stale benchmark build metadata ({fields} changed): {metadata}") + + binaries = record.get("binaries", {}) + resolved: dict[str, str] = {} + for suite in args.suite: + stored = binaries.get(suite) + if stored is None: + raise ValueError(f"benchmark suite {suite!r} was not recorded in {metadata}") + path = Path(stored["path"]) + if not path.is_file(): + raise ValueError(f"recorded benchmark binary does not exist: {path}") + current = { + "path": str(path.resolve()), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + if current != stored: + raise ValueError(f"recorded benchmark binary changed: {path}") + resolved[suite] = str(path.resolve()) + + return resolved + + +def validate_build_record(args: argparse.Namespace) -> None: + for suite, path in validated_build_binaries(args).items(): + print(f"{suite}={path}") + + +def write_manifest(args: argparse.Namespace) -> None: + settings = dict(setting.split("=", 1) for setting in args.setting) + manifest = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "settings": settings, + "suites": args.suite, + "filters": args.filter, + "machine_record": str(Path(args.machine_record).resolve()), + "baseline": revision_record( + Path(args.baseline_worktree), + Path(args.baseline_target), + args.baseline_binary, + ), + "candidate": revision_record( + Path(args.candidate_worktree), + Path(args.candidate_target), + args.candidate_binary, + ), + } + output = Path(args.output) + output.write_text(f"{json.dumps(manifest, indent=2, sort_keys=True)}\n", encoding="utf-8") + + +def timing_ns(field: str) -> float: + match = TIMING.search(field.strip()) + if match is None: + raise ValueError(f"cannot parse Divan timing from {field!r}") + + return float(match.group("value")) * UNIT_TO_NS[match.group("unit")] + + +def parse_divan(path: Path) -> dict[str, float]: + """Return benchmark paths and median nanoseconds from one Divan table.""" + + parents: dict[int, str] = {} + timings: dict[str, float] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + fields = re.split(r"\s+│\s+", line) + tree_match = TREE_ROW.match(fields[0]) + if tree_match is None: + continue + + depth = len(tree_match.group("prefix")) // 3 + body = tree_match.group("body").rstrip() + timing_match = TIMING.search(body) + name = body[: timing_match.start()].rstrip() if timing_match else body.strip() + parents = {level: parent for level, parent in parents.items() if level < depth} + + if timing_match is None: + parents[depth] = name + continue + if len(fields) < 3: + raise ValueError(f"timed Divan row has no median column in {path}: {line}") + + components = [parents[level] for level in sorted(parents) if level < depth] + benchmark = "/".join([*components, name]) + if benchmark in timings: + raise ValueError(f"duplicate benchmark {benchmark!r} in {path}") + timings[benchmark] = timing_ns(fields[2]) + + if not timings: + raise ValueError(f"no Divan benchmark timings found in {path}") + + return timings + + +def read_measurements(directory: Path) -> dict[tuple[str, str, int, str], float]: + measurements: dict[tuple[str, str, int, str], float] = {} + for path in sorted(directory.glob("*.txt")): + match = RESULT_FILE.match(path.name) + if match is None: + continue + suite = match.group("suite") + revision = match.group("revision") + pair = int(match.group("pair")) + for benchmark, median_ns in parse_divan(path).items(): + measurements[suite, revision, pair, benchmark] = median_ns + + if not measurements: + raise ValueError(f"no measured result files found in {directory}") + + return measurements + + +def summarize(measurements: dict[tuple[str, str, int, str], float]) -> list[BenchmarkSummary]: + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + suites = {suite for suite, _ in inventories} + comparable = { + (suite, benchmark) + for suite in suites + for benchmark in inventories.get((suite, "baseline"), set()) & inventories.get((suite, "candidate"), set()) + } + groups = { + (suite, pair, benchmark) for suite, _, pair, benchmark in measurements if (suite, benchmark) in comparable + } + incomplete = [ + group + for group in groups + if (group[0], "baseline", group[1], group[2]) not in measurements + or (group[0], "candidate", group[1], group[2]) not in measurements + ] + if incomplete: + raise ValueError(f"unpaired benchmark measurements: {sorted(incomplete)!r}") + if not groups: + raise ValueError("unpaired benchmark measurements: no comparable benchmarks") + + by_benchmark: dict[tuple[str, str], list[tuple[float, float]]] = {} + for suite, pair, benchmark in sorted(groups): + baseline = measurements[suite, "baseline", pair, benchmark] + candidate = measurements[suite, "candidate", pair, benchmark] + by_benchmark.setdefault((suite, benchmark), []).append((baseline, candidate)) + + summaries = [] + for (suite, benchmark), pairs in sorted(by_benchmark.items()): + baseline_values = [baseline for baseline, _ in pairs] + candidate_values = [candidate for _, candidate in pairs] + ratios = [candidate / baseline for baseline, candidate in pairs] + median_ratio = statistics.median(ratios) + summaries.append( + BenchmarkSummary( + suite=suite, + benchmark=benchmark, + pairs=len(pairs), + baseline_median_ns=statistics.median(baseline_values), + candidate_median_ns=statistics.median(candidate_values), + median_ratio=median_ratio, + minimum_ratio=min(ratios), + maximum_ratio=max(ratios), + ratio_mad=statistics.median(abs(ratio - median_ratio) for ratio in ratios), + ) + ) + + return summaries + + +def inventory_differences( + measurements: dict[tuple[str, str, int, str], float], +) -> list[tuple[str, str, str]]: + """Return benchmarks that exist in only one revision.""" + + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + differences = [] + for suite in sorted({suite for suite, _ in inventories}): + baseline = inventories.get((suite, "baseline"), set()) + candidate = inventories.get((suite, "candidate"), set()) + differences.extend((suite, "baseline only", benchmark) for benchmark in baseline - candidate) + differences.extend((suite, "candidate only", benchmark) for benchmark in candidate - baseline) + + return sorted(differences) + + +def format_ns(value: float) -> str: + for divisor, unit in ((1_000_000_000, "s"), (1_000_000, "ms"), (1_000, "µs")): + if value >= divisor: + return f"{value / divisor:.3f} {unit}" + + return f"{value:.3f} ns" + + +def write_summary( + output_directory: Path, + summaries: list[BenchmarkSummary], + differences: Iterable[tuple[str, str, str]] = (), +) -> None: + csv_path = output_directory / "ratios.csv" + with csv_path.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=list(asdict(summaries[0]))) + writer.writeheader() + writer.writerows(asdict(summary) for summary in summaries) + + markdown = [ + "# RowFn benchmark comparison", + "", + "Ratios are paired candidate/baseline medians. Lower is faster.", + "", + "| Suite | Benchmark | Pairs | Baseline | Candidate | Ratio | Change | MAD |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for summary in sorted(summaries, key=lambda result: result.median_ratio, reverse=True): + change = (summary.median_ratio - 1.0) * 100.0 + markdown.append( + f"| {summary.suite} | `{summary.benchmark}` | {summary.pairs} " + f"| {format_ns(summary.baseline_median_ns)} " + f"| {format_ns(summary.candidate_median_ns)} " + f"| {summary.median_ratio:.6f} | {change:+.2f}% | {summary.ratio_mad:.6f} |" + ) + differences = list(differences) + if differences: + markdown.extend( + [ + "", + "## Unpaired benchmark inventory", + "", + "These benchmarks were recorded for only one revision and are excluded from ratios.", + "", + ] + ) + markdown.extend(f"- `{suite}/{benchmark}`: {revision}." for suite, revision, benchmark in differences) + markdown.append("") + (output_directory / "summary.md").write_text("\n".join(markdown), encoding="utf-8") + + +def summarize_directory(args: argparse.Namespace) -> None: + output_directory = Path(args.output_directory) + measurements = read_measurements(output_directory / "measured") + summaries = summarize(measurements) + write_summary(output_directory, summaries, inventory_differences(measurements)) + + +def argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(required=True) + + manifest = subparsers.add_parser("manifest", help="capture revisions and executable hashes") + manifest.add_argument("--output", required=True) + manifest.add_argument("--machine-record", required=True) + manifest.add_argument("--baseline-worktree", required=True) + manifest.add_argument("--candidate-worktree", required=True) + manifest.add_argument("--baseline-target", required=True) + manifest.add_argument("--candidate-target", required=True) + manifest.add_argument("--setting", action="append", default=[]) + manifest.add_argument("--suite", action="append", default=[]) + manifest.add_argument("--filter", action="append", default=[]) + manifest.add_argument("--baseline-binary", action="append", default=[]) + manifest.add_argument("--candidate-binary", action="append", default=[]) + manifest.set_defaults(function=write_manifest) + + record_build = subparsers.add_parser("record-build", help="record reusable benchmark binaries") + record_build.add_argument("--output", required=True) + record_build.add_argument("--worktree", required=True) + record_build.add_argument("--target", required=True) + record_build.add_argument("--setting", action="append", default=[]) + record_build.add_argument("--binary", action="append", default=[]) + record_build.set_defaults(function=write_build_record) + + validate_build = subparsers.add_parser("validate-build", help="validate a reusable build") + validate_build.add_argument("--metadata", required=True) + validate_build.add_argument("--worktree", required=True) + validate_build.add_argument("--target", required=True) + validate_build.add_argument("--setting", action="append", default=[]) + validate_build.add_argument("--suite", action="append", default=[]) + validate_build.set_defaults(function=validate_build_record) + + summary = subparsers.add_parser("summarize", help="write ratios.csv and summary.md") + summary.add_argument("output_directory") + summary.set_defaults(function=summarize_directory) + + return parser + + +def main() -> None: + parser = argument_parser() + args = parser.parse_args() + try: + args.function(args) + except (OSError, subprocess.CalledProcessError, ValueError) as error: + parser.error(str(error)) + + +if __name__ == "__main__": + main() diff --git a/scripts/tests/test_rowfn_benchmark.py b/scripts/tests/test_rowfn_benchmark.py new file mode 100644 index 00000000000..ace0797c6a3 --- /dev/null +++ b/scripts/tests/test_rowfn_benchmark.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +REPO_ROOT = Path(__file__).resolve().parents[2] +SCRIPT = REPO_ROOT / "scripts" / "rowfn_benchmark.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("rowfn_benchmark", SCRIPT) + assert spec is not None + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def write_divan(path: Path, rows: list[str]) -> None: + path.write_text( + "\n".join( + [ + "Timer precision: 20 ns", + "bench fastest │ slowest │ median │ mean │ samples │ iters", + *rows, + "", + ] + ), + encoding="utf-8", + ) + + +class RowFnBenchmarkTest(unittest.TestCase): + def setUp(self) -> None: + self.module = load_module() + self.temporary_directory = tempfile.TemporaryDirectory() + self.directory = Path(self.temporary_directory.name) + + def tearDown(self) -> None: + self.temporary_directory.cleanup() + + def test_parse_divan_preserves_nested_benchmark_names_and_converts_units(self) -> None: + output = self.directory / "result.txt" + write_divan( + output, + [ + "├─ non_nullable │ │ │ │ │", + "│ ├─ 2 17.18 µs │ 18 µs │ 17.33 µs │ 17.4 µs │ 100 │ 100", + "│ ╰─ 32 6.709 µs │ 8 µs │ 6.829 µs │ 7 µs │ 100 │ 100", + "╰─ nullable │ │ │ │ │", + " ╰─ 2 799.7 ns │ 1 µs │ 979.7 ns │ 986 ns │ 100 │ 100", + ], + ) + + self.assertEqual( + self.module.parse_divan(output), + { + "non_nullable/2": 17_330.0, + "non_nullable/32": 6_829.0, + "nullable/2": 979.7, + }, + ) + + def test_summarize_writes_paired_ratios_and_slowest_first_markdown(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": ["├─ add 12 ns │ 12 ns │ 12 ns │ 12 ns │ 100 │ 100"], + "numeric-baseline-2.txt": ["├─ add 20 ns │ 20 ns │ 20 ns │ 20 ns │ 100 │ 100"], + "numeric-candidate-2.txt": ["├─ add 18 ns │ 18 ns │ 18 ns │ 18 ns │ 100 │ 100"], + "numeric-baseline-3.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-3.txt": ["├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100"], + "numeric-baseline-4.txt": ["├─ mul 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-4.txt": ["├─ mul 9 ns │ 9 ns │ 9 ns │ 9 ns │ 100 │ 100"], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + summaries = self.module.summarize(self.module.read_measurements(measured)) + self.module.write_summary(self.directory, summaries) + + add = next(summary for summary in summaries if summary.benchmark == "add") + self.assertEqual(add.pairs, 3) + self.assertAlmostEqual(add.median_ratio, 1.1) + self.assertAlmostEqual(add.ratio_mad, 0.1) + + csv_output = (self.directory / "ratios.csv").read_text(encoding="utf-8") + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("suite,benchmark,pairs", csv_output) + self.assertLess(markdown.index("`add`"), markdown.index("`mul`")) + + def test_summarize_rejects_unpaired_measurements(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + write_divan( + measured / "numeric-baseline-1.txt", + ["╰─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + ) + + with self.assertRaisesRegex(ValueError, "unpaired benchmark measurements"): + self.module.summarize(self.module.read_measurements(measured)) + + def test_summarize_excludes_and_reports_revision_only_benchmarks(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": [ + "├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100", + "╰─ candidate 5 ns │ 5 ns │ 5 ns │ 5 ns │ 100 │ 100", + ], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + measurements = self.module.read_measurements(measured) + summaries = self.module.summarize(measurements) + differences = self.module.inventory_differences(measurements) + self.module.write_summary(self.directory, summaries, differences) + + self.assertEqual([summary.benchmark for summary in summaries], ["add"]) + self.assertEqual(differences, [("numeric", "candidate only", "candidate")]) + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("`numeric/candidate`: candidate only.", markdown) + + def test_build_record_validates_identity_and_executable(self) -> None: + target = self.directory / "target" + target.mkdir() + binary = target / "binary_ops-123" + binary.write_bytes(b"first binary") + metadata = target / "rowfn-benchmark-build.json" + identity = { + "settings": {"codegen_units": "1", "lto": "fat"}, + "toolchain": {"rustc": "rustc 1.97.1", "cargo": "cargo 1.97.1"}, + "revision": {"head": "abc123", "dirty_state_sha256": "clean"}, + } + arguments = SimpleNamespace( + output=str(metadata), + worktree=str(self.directory), + target=str(target), + setting=["codegen_units=1", "lto=fat"], + binary=[f"numeric={binary}"], + ) + + with mock.patch.object(self.module, "build_identity", return_value=identity): + self.module.write_build_record(arguments) + + validation = SimpleNamespace( + metadata=str(metadata), + worktree=str(self.directory), + target=str(target), + setting=["codegen_units=1", "lto=fat"], + suite=["numeric"], + ) + with mock.patch.object(self.module, "build_identity", return_value=identity): + self.assertEqual( + self.module.validated_build_binaries(validation), + {"numeric": str(binary.resolve())}, + ) + + changed_identities = { + "settings": {**identity, "settings": {"codegen_units": "16", "lto": "false"}}, + "toolchain": { + **identity, + "toolchain": {"rustc": "rustc 1.98.0", "cargo": "cargo 1.98.0"}, + }, + "revision": { + **identity, + "revision": {"head": "def456", "dirty_state_sha256": "changed"}, + }, + } + for field, changed_identity in changed_identities.items(): + with ( + self.subTest(field=field), + mock.patch.object(self.module, "build_identity", return_value=changed_identity), + self.assertRaisesRegex(ValueError, f"{field} changed"), + ): + self.module.validated_build_binaries(validation) + + binary.write_bytes(b"second binary") + with ( + mock.patch.object(self.module, "build_identity", return_value=identity), + self.assertRaisesRegex(ValueError, "binary changed"), + ): + self.module.validated_build_binaries(validation) + + +if __name__ == "__main__": + unittest.main() diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 2af2eacf238..68bd189ef11 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -134,6 +134,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_executor" +harness = false + [[bench]] name = "kleene_bool" harness = false @@ -213,6 +217,10 @@ harness = false name = "validity_is_valid" harness = false +[[bench]] +name = "strict_validity" +harness = false + [[bench]] name = "dict_unreferenced_mask" harness = false diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 6a07d03f50b..ccd440ff939 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -38,10 +38,60 @@ static SESSION: LazyLock = LazyLock::new(array_session); const LEN: usize = 32_768; +const ROWFN_MATRIX_CASES: &[(usize, RowFnShape)] = &[ + (128, RowFnShape::PerRowPerRow), + (128, RowFnShape::PerRowConstant), + (128, RowFnShape::ConstantPerRow), + (128, RowFnShape::PerRowNullableConstant), + (LEN, RowFnShape::PerRowPerRow), + (LEN, RowFnShape::PerRowConstant), + (LEN, RowFnShape::ConstantPerRow), + (LEN, RowFnShape::PerRowNullableConstant), +]; + +#[derive(Clone, Copy, Debug)] +enum RowFnShape { + PerRowPerRow, + PerRowConstant, + ConstantPerRow, + PerRowNullableConstant, +} + /// Decimal Mul and Div cost far more per lane than Add, so they run over a shorter array to keep /// the instrumented CodSpeed runs quick. const DECIMAL_MUL_DIV_LEN: usize = 8_192; +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_add(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Add); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_subtract(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Sub); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_multiply(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Mul); +} + +fn bench_rowfn_shape(bencher: Bencher, len: usize, shape: RowFnShape, operator: Operator) { + let per_row = + || PrimitiveArray::from_iter((0..len).map(|index| (index % 1_024) as i64 + 1)).into_array(); + let constant = || ConstantArray::new(17_i64, len).into_array(); + let nullable_constant = || ConstantArray::new(Some(17_i64), len).into_array(); + + let (lhs, rhs) = match shape { + RowFnShape::PerRowPerRow => (per_row(), per_row()), + RowFnShape::PerRowConstant => (per_row(), constant()), + RowFnShape::ConstantPerRow => (constant(), per_row()), + RowFnShape::PerRowNullableConstant => (per_row(), nullable_constant()), + }; + + bench_primitive(bencher, lhs, rhs, operator); +} + #[divan::bench] fn add_i64_nonnull(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); @@ -170,6 +220,14 @@ fn div_i64_nonnull(bencher: Bencher) { bench_primitive(bencher, lhs, rhs, Operator::Div); } +#[divan::bench] +fn div_i64_nullable(bencher: Bencher) { + let lhs = primitive_nullable(1_000_000, 7).into_array(); + let rhs = primitive_nullable(17, 5).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn sub_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 4a399760dc2..9e6dd3e4e5b 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -4,6 +4,7 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; @@ -38,6 +39,7 @@ const ARRAY_SIZE: usize = 65_536; fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); bencher + .counter(ItemsCount::new(ARRAY_SIZE)) .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input @@ -49,6 +51,31 @@ fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { }); } +fn u8_array(offset: u8) -> ArrayRef { + (0u8..=u8::MAX) + .cycle() + .take(ARRAY_SIZE) + .map(|value| value.wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn i32_array(offset: i32) -> ArrayRef { + (0i32..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn u64_array(offset: u64) -> ArrayRef { + (0u64..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + fn bool_array(rng: &mut StdRng) -> ArrayRef { BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() } @@ -87,6 +114,13 @@ fn float_array(rng: &mut StdRng) -> ArrayRef { .into_array() } +fn f32_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f32..1.0)) + .collect::>() + .into_array() +} + fn string_array(rng: &mut StdRng) -> ArrayRef { VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { let len = rng.random_range(1usize..24); @@ -153,6 +187,14 @@ fn compare_int_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Gte); } +#[divan::bench] +fn compare_int_constant_lhs(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + let arr = int_array(&mut rng); + bench_compare(bencher, constant, arr, Operator::Gte); +} + #[divan::bench] fn compare_int_eq(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -161,6 +203,55 @@ fn compare_int_eq(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Eq); } +#[divan::bench] +fn compare_i32(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = i32_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_i32_constant(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = ConstantArray::new(1_000_000i32, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = u8_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8_constant(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = ConstantArray::new(127u8, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_constant(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = ConstantArray::new(1_000_000u64, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_eq(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Eq); +} + #[divan::bench] fn compare_float(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -169,6 +260,30 @@ fn compare_float(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[divan::bench] +fn compare_float_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_f32_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + #[divan::bench] fn compare_decimal(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs new file mode 100644 index 00000000000..84a2e029412 --- /dev/null +++ b/vortex-array/benches/row_fn_executor.rs @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares owned-output, sink-writing, and hand-written primitive row loops. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::OutputSink; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const ROWS: usize = 65_536; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +#[derive(Clone)] +struct RowWrappingAdd; + +impl RowFn for RowWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_wrapping_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) + } +} + +#[derive(Clone)] +struct RowCheckedAdd; + +impl RowFn for RowCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_checked_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(i64, i64), i64, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(checked_add_error()); + } + Ok(()) + }, + ) + } +} + +/// Keep error construction out of the benchmarked success path. +#[cold] +#[inline(never)] +fn checked_add_error() -> VortexError { + vortex_err!("integer overflow in row checked add") +} + +/// A benchmark sink that writes one `i64` per row. +struct I64Sink( + /// The output values written by the row loop. + BufferMut, +); + +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +#[derive(Clone)] +struct RowSinkWrappingAdd; + +impl RowFn for RowSinkWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_wrapping_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), I64Sink, _>(|(lhs, rhs), out| { + *out = lhs.wrapping_add(rhs); + }) + } +} + +#[derive(Clone)] +struct RowSinkCheckedAdd; + +impl RowFn for RowSinkCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_checked_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), UninitElementSink, _>( + |(lhs, rhs), output| -> VortexResult { + let value = lhs.checked_add(rhs).ok_or_else(checked_add_error)?; + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }, + ) + } +} + +fn inputs() -> (ArrayRef, ArrayRef) { + let lhs = (0..ROWS) + .map(|index| index as i64) + .collect::>() + .into_array(); + let rhs = (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>() + .into_array(); + (lhs, rhs) +} + +fn constant_inputs() -> (ArrayRef, ArrayRef) { + let (lhs, _) = inputs(); + let rhs = ConstantArray::new(Scalar::from(7i64), ROWS).into_array(); + (lhs, rhs) +} + +fn nullable_inputs() -> (ArrayRef, ArrayRef) { + let lhs = PrimitiveArray::new( + (0..ROWS).map(|index| index as i64).collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(5))), + ) + .into_array(); + let rhs = PrimitiveArray::new( + (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(7))), + ) + .into_array(); + (lhs, rhs) +} + +fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) +where + F: RowFn, +{ + bencher + .with_inputs(make_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + function + .clone() + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, nullable_inputs); +} + +#[divan::bench] +fn handrolled_sink_wrapping_add(bencher: Bencher) { + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = lhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let rhs = rhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let mut output = BufferMut::zeroed(ROWS); + for ((out, lhs), rhs) in output + .as_mut_slice() + .iter_mut() + .zip(lhs.as_slice()) + .zip(rhs.as_slice()) + { + *out = lhs.wrapping_add(*rhs); + } + PrimitiveArray::new(output.freeze(), Validity::NonNullable).into_array() + }); +} + +#[divan::bench] +fn row_checked_add(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, inputs); +} + +#[divan::bench] +fn row_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, nullable_inputs); +} + +#[divan::bench] +fn row_sink_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowSinkCheckedAdd, nullable_inputs); +} + +#[divan::bench] +fn row_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, nullable_inputs); +} diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs new file mode 100644 index 00000000000..9a713220fab --- /dev/null +++ b/vortex-array/benches/strict_validity.rs @@ -0,0 +1,216 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks how the row lifting applies validity to a dense kernel. +//! +//! Both arms run the same kernel over the same nullable input and differ only in how the output's +//! nulls are restored: +//! +//! - `lazy` is what the lifting behind [`RowFn`] does: conjoin the input validities (itself lazy) +//! and hand the resulting boolean array straight to `mask`. +//! - `eager` is the strategy it replaced: materialize the conjunction into a `Mask`, copy it into a +//! `BitBuffer`, wrap that in a `BoolArray`, and mask with it. +//! +//! `chain` composes three of the same function, which is where a per-call materialization +//! compounds. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const SIZES: &[usize] = &[65_536, 1 << 20]; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +/// The shared kernel: double every lane, ignoring validity. +fn doubled(input: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let input = input.clone().execute::(ctx)?; + let values: Buffer = input + .as_slice::() + .iter() + .map(|value| value.wrapping_mul(2)) + .collect(); + Ok(PrimitiveArray::new(values, Validity::NonNullable).into_array()) +} + +/// Dense row function: the lifting decides how validity is applied. +/// +/// Its [`reduce_encoded`](RowFn::reduce_encoded) answers with [`doubled`] before the row loop is +/// reached, so this arm runs exactly the kernel [`EagerDouble`] runs and the two differ only in +/// validity. The row closure below is what makes it a [`RowFn`] at all, and never executes. +#[derive(Clone)] +struct LazyDouble; + +impl RowFn for LazyDouble { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.lazy_double"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i32,), i32>(|(value,)| value.wrapping_mul(2)) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + doubled(&args[0], ctx).map(|output| Some(RowExecution::Output(output))) + } +} + +/// The same function, applying validity the way the adapter used to: materialize a mask first. +#[derive(Clone)] +struct EagerDouble; + +impl ScalarFnVTable for EagerDouble { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.eager_double"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + ChildName::from("input") + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::I32, args[0].nullability())) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let valid = input.validity()?.execute_mask(args.row_count(), ctx)?; + let values = doubled(&input, ctx)?; + + if valid.all_true() { + return values.cast(DType::Primitive(PType::I32, input.dtype().nullability())); + } + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + values.mask(mask) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// A nullable i32 column with array-backed validity (~10% nulls). +fn nullable_input(len: usize) -> ArrayRef { + PrimitiveArray::new( + (0..len as i32).collect::>(), + Validity::from_iter((0..len).map(|index| !index.is_multiple_of(10))), + ) + .into_array() +} + +fn bench_depth(bencher: Bencher, function: F, len: usize, depth: usize) +where + F: ScalarFnVTable + Clone, +{ + bencher + .with_inputs(|| nullable_input(len)) + .bench_local_values(|input| { + let mut ctx = SESSION.create_execution_ctx(); + let mut array = input; + for _ in 0..depth { + array = function + .clone() + .try_new_array(len, EmptyOptions, [array]) + .unwrap() + .into_array(); + } + array.execute::(&mut ctx).unwrap() + }); +} + +#[divan::bench(args = SIZES)] +fn lazy(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn eager(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn lazy_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 3); +} + +#[divan::bench(args = SIZES)] +fn eager_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 3); +} diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index d25a652ee57..a36d0a22bde 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -4,9 +4,9 @@ //! Native comparison kernels. //! //! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every -//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from -//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise -//! comparator for nested types. There is no Arrow fallback. +//! comparison directly over Vortex canonical arrays: bit buffers for booleans, row or fused lane +//! kernels for primitives, lane kernels for decimals, binary views for strings and bytes, and a +//! row-wise comparator for nested types. There is no Arrow fallback. //! //! Floating point values compare with Vortex's total ordering (`NaN` is the largest value, //! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics. @@ -211,7 +211,7 @@ fn compare_arrays( ) .into_array()), DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx), - DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), + DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, ctx), DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 1247358dce1..93afcf538ed 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,28 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays via bit-packing lane kernels. +//! Primitive comparison execution through [`RowFn`]. + +mod columnar; +mod operand; -use vortex_buffer::BitBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_bail; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; +#[cfg(target_arch = "x86_64")] +use crate::arrays::Constant; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::compare::collect_bits; -use crate::scalar_fn::fns::binary::compare::collect_zip_bits; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; /// Compare two primitive arrays of the same [`PType`]. /// @@ -32,99 +34,125 @@ pub(super) fn compare_primitive( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - match_each_native_ptype!(ptype, |T| { - compare_primitive_typed::(lhs, rhs, op, nullability, ctx) - }) + compare_primitive_with_path(lhs, rhs, op, PrimitiveComparisonPath::Auto, ctx) +} + +/// Selects automatic production dispatch or a forced implementation in tests. +#[derive(Clone, Copy)] +pub(super) enum PrimitiveComparisonPath { + /// Use the architecture and operand-specific production policy. + Auto, + + /// Force row execution. + #[cfg(test)] + Row, + + /// Force fused columnar execution. + #[cfg(test)] + Columnar, } -fn compare_primitive_typed( +/// Compare primitives through the selected implementation. +pub(super) fn compare_primitive_with_path( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, + path: PrimitiveComparisonPath, ctx: &mut ExecutionCtx, ) -> VortexResult { - let len = lhs.len(); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - if lhs.len() != rhs.len() { - vortex_bail!( - "compare operator requires equal lengths, got {} and {}", - lhs.len(), - rhs.len() - ); - } - - let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; - - let bits = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slices(lhs, rhs, op), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => compare_slice_constant(lhs, *rhs, op), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slice_constant(rhs, *lhs, op.swap()), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Unreachable through `execute_compare` (constant-constant is folded there), but - // cheap to answer anyway. - BitBuffer::full(apply_op(*lhs, *rhs, op), len) - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { - return Ok( - ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) - .into_array(), - ); + let use_columnar = match path { + PrimitiveComparisonPath::Auto => { + #[cfg(target_arch = "x86_64")] + { + use_columnar_comparison(lhs, rhs, op)? + } + #[cfg(not(target_arch = "x86_64"))] + { + false + } } + #[cfg(test)] + PrimitiveComparisonPath::Row => false, + #[cfg(test)] + PrimitiveComparisonPath::Columnar => true, }; + if use_columnar { + return columnar::compare_primitive(lhs, rhs, op, ctx); + } - Ok(BoolArray::try_new(bits, validity)?.into_array()) + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + execute_rows(&PrimitiveCompare, &op, &args, ctx) } -#[inline(always)] -fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { - match op { - CompareOperator::Eq => lhs.is_eq(rhs), - CompareOperator::NotEq => !lhs.is_eq(rhs), - CompareOperator::Gt => lhs.is_gt(rhs), - CompareOperator::Gte => lhs.is_ge(rhs), - CompareOperator::Lt => lhs.is_lt(rhs), - CompareOperator::Lte => lhs.is_le(rhs), +/// Internal row execution for primitive comparison operators. +#[derive(Clone)] +struct PrimitiveCompare; + +impl RowFn for PrimitiveCompare { + type Options = CompareOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + // `PrimitiveCompare` is a private implementation detail of `Binary`: it is never registered + // or serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. + ScalarFnVTable::id(&Binary) } -} -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { - // Dispatch the operator outside the lane loop so each instantiation vectorizes a single - // branch-free predicate. - match op { - CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = + PType::try_from(args.first().ok_or_else(|| { + vortex_err!("a comparison operator takes two operands, got none") + })?)?; + + match_each_native_ptype!(ptype, |T| { visit_compare::(*op, visitor) }) } } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { +#[cfg(target_arch = "x86_64")] +fn use_columnar_comparison( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, +) -> VortexResult { + let ptype = PType::try_from(lhs.dtype())?; + Ok(match (ptype, op) { + // Equality bit-packs efficiently for every type supported by the columnar path. + (PType::I64 | PType::U64 | PType::F64, CompareOperator::Eq | CompareOperator::NotEq) => { + true + } + // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit + // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. + (PType::I64 | PType::F64, _) => true, + // LLVM 22 vectorizes the mixed-constant RowFn loop at 16 CGUs without LTO. However, the + // fused comparison and bit-packing path is still about 38% faster in + // `compare_u64_constant`. Recheck that benchmark before changing this dispatch. + (PType::U64, _) => lhs.is::() || rhs.is::(), + _ => false, + }) +} + +fn visit_compare(op: CompareOperator, visitor: V) -> VortexResult +where + T: NativePType, + V: RowVisitor, +{ match op { - CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), - CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), - CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), - CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), - CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), - CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), + CompareOperator::Eq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)), + CompareOperator::NotEq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| !lhs.is_eq(rhs)), + CompareOperator::Gt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_gt(rhs)), + CompareOperator::Gte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_ge(rhs)), + CompareOperator::Lt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs)), + CompareOperator::Lte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_le(rhs)), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs new file mode 100644 index 00000000000..6728437e6a8 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fused comparison and bit-packing for wide primitive lanes. +//! +//! Production uses this implementation only for measured x86 paths. Keeping it portable lets the +//! semantic tests exercise the RowFn and fused paths on every target. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use super::operand::PrimitiveOperand; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare primitive operands with one fused comparison and bit-packing loop. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match PType::try_from(lhs.dtype())? { + PType::I64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::U64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::F64 => compare_primitive_typed::(lhs, rhs, op, ctx), + ptype => vortex_bail!("columnar comparison is not selected for {ptype}"), + } +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + let bits = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slices(lhs, rhs, op), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => compare_slice_constant(lhs, *rhs, op), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slice_constant(rhs, *lhs, op.swap()), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => BitBuffer::full(apply_op(*lhs, *rhs, op), len), + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { + match op { + CompareOperator::Eq => lhs.is_eq(rhs), + CompareOperator::NotEq => !lhs.is_eq(rhs), + CompareOperator::Gt => lhs.is_gt(rhs), + CompareOperator::Gte => lhs.is_ge(rhs), + CompareOperator::Lt => lhs.is_lt(rhs), + CompareOperator::Lte => lhs.is_le(rhs), + } +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + } +} + +fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |lhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |lhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |lhs: T| lhs.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |lhs: T| lhs.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |lhs: T| lhs.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |lhs: T| lhs.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs similarity index 78% rename from vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs rename to vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs index 71d1122fc79..1563b8e68e6 100644 --- a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Decoding shared by primitive binary operators. +//! Operand decoding for the fused primitive comparison path. use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -15,19 +15,33 @@ use crate::validity::Validity; /// A materialized primitive column, a non-null constant, or an all-null constant. pub(super) enum PrimitiveOperand { + /// A per-row primitive column and its validity. Array { + /// The materialized values. values: Buffer, + + /// The validity of the values. validity: Validity, }, + + /// A non-null value repeated for every row. Constant { + /// The repeated value. value: T, + + /// The number of repeated rows. len: usize, + + /// The validity implied by the constant's dtype. validity: Validity, }, + + /// An all-null constant with this row count. Null(usize), } impl PrimitiveOperand { + /// Decode an operand once for the fused comparison loop. pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Some(constant) = array.as_opt::() { return Ok( @@ -49,9 +63,11 @@ impl PrimitiveOperand { let array = array.clone().execute::(ctx)?; let validity = array.validity()?; let values = array.into_buffer::(); + Ok(Self::Array { values, validity }) } + /// Return the logical row count. pub(super) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), @@ -59,6 +75,7 @@ impl PrimitiveOperand { } } + /// Return the operand validity. pub(super) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index 9831a963354..d0f7a9b5e57 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -12,7 +12,9 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::VTable; use crate::array_session; +use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; @@ -21,6 +23,7 @@ use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::ScalarFn; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; @@ -39,6 +42,8 @@ use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::primitive::PrimitiveComparisonPath; +use crate::scalar_fn::fns::binary::compare::primitive::compare_primitive_with_path; use crate::scalar_fn::fns::binary::scalar_cmp; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -429,6 +434,142 @@ fn float_total_order() { ); } +#[rstest] +#[case::row_eq(PrimitiveComparisonPath::Row, CompareOperator::Eq)] +#[case::row_not_eq(PrimitiveComparisonPath::Row, CompareOperator::NotEq)] +#[case::row_lt(PrimitiveComparisonPath::Row, CompareOperator::Lt)] +#[case::columnar_eq(PrimitiveComparisonPath::Columnar, CompareOperator::Eq)] +#[case::columnar_not_eq(PrimitiveComparisonPath::Columnar, CompareOperator::NotEq)] +#[case::columnar_lt(PrimitiveComparisonPath::Columnar, CompareOperator::Lt)] +fn test_primitive_comparison_paths_preserve_semantics_and_encoding( + #[case] path: PrimitiveComparisonPath, + #[case] op: CompareOperator, +) -> VortexResult<()> { + let lhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::NAN, // Null on the left. + -0.0, // Signed zero ordering. + 1.0, // A finite value below NaN. + f64::NAN, // Null on the right. + ], + Validity::from_iter([ + true, // + false, // + true, // + true, // + true, // + ]), + ) + .into_array(); + let rhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::INFINITY, // Null on the left. + 0.0, // Signed zero ordering. + f64::NAN, // A finite value below NaN. + 1.0, // Null on the right. + ], + Validity::from_iter([ + true, // + true, // + true, // + true, // + false, // + ]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let actual = compare_primitive_with_path(&lhs, &rhs, op, path, &mut ctx)?; + let expected = match op { + CompareOperator::Eq => [ + Some(true), // Equal NaNs. + None, // Null on the left. + Some(false), // Distinct signed zeroes. + Some(false), // A finite value and NaN. + None, // Null on the right. + ], + CompareOperator::NotEq | CompareOperator::Lt => [ + Some(false), // Equal NaNs. + None, // Null on the left. + Some(true), // Distinct signed zeroes. + Some(true), // A finite value and NaN. + None, // Null on the right. + ], + _ => unreachable!(), + }; + let expected = BoolArray::from_iter(expected); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); + + // This encoding difference is intentional: the fused path materializes bits and validity + // together, while the RowFn path keeps masking lazy. + match path { + PrimitiveComparisonPath::Columnar => assert_eq!(actual.encoding_id(), Bool.id()), + PrimitiveComparisonPath::Row => assert!(actual.as_opt::().is_some()), + PrimitiveComparisonPath::Auto => unreachable!(), + } + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[rstest] +#[case::i64_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::i64_not_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::u64_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::u64_not_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::f64_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::f64_not_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +fn test_primitive_equality_auto_uses_columnar_for_supported_ptype( + #[case] lhs: ArrayRef, + #[case] rhs: ArrayRef, + #[case] op: CompareOperator, + #[case] expected: [bool; 3], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let actual = + compare_primitive_with_path(&lhs, &rhs, op, PrimitiveComparisonPath::Auto, &mut ctx)?; + + assert_eq!(actual.encoding_id(), Bool.id()); + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + #[rstest] #[case(Operator::Eq, [true, false, true, true])] #[case(Operator::Lt, [false, true, false, false])] diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 80faff20e0a..a5b9fe70539 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -43,7 +43,6 @@ mod compare; pub use compare::*; mod numeric; pub(crate) use numeric::*; -mod primitive_operand; use crate::scalar::NumericOperator; use crate::scalar::Scalar; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..054846b7ef7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure evidence without +//! scanning the finished output. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,33 +15,18 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. +/// Keep this wrapper inlineable so captured constants can become loop invariants in the caller. +/// The lane kernels retain their own inlining decisions. #[inline] pub(super) fn checked_lanes( source: S, @@ -61,7 +48,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +59,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - mut apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6dc0de0fbea..a0c427b142e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,16 +4,19 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a +//! [`RowFn`](crate::scalar_fn::unstable::row::RowFn), which owns null handling, constants, and +//! validity for them; see [`row`]. Decimal keeps its own columnar implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; -#[cfg(test)] -mod tests; +mod row; use decimal::execute_numeric_decimal; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -81,3 +84,6 @@ fn build_empty_result( Ok(Canonical::empty(&result_dtype).into_array()) } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 8fd53d15216..7cbe88c7c8e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,73 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! Checked arithmetic for one primitive row. + +use std::ops::BitOrAssign; + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::validity::Validity; - -struct CheckedAdd; -struct CheckedSub; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedMul; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedDiv; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. - const ERROR: &'static str; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; +/// OR-reducible evidence that a row failed, with [`Default`] meaning success. +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - /// How this operation reports a failing lane. See [`Failure`]. - type Failure: Failure; +impl Failure for bool {} +impl Failure for u8 {} +impl Failure for u16 {} +impl Failure for u32 {} +impl Failure for u64 {} - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. - fn apply(lhs: T, rhs: T) -> (T, Self::Failure); +/// One arithmetic operator at one width, split into its value and failure evidence. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. + const ERROR: &'static str; - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); + /// How this operation reports a failing row. See [`Failure`]. + type Fail: Failure; - (failed == Self::Failure::default()).then_some(value) - } + /// The result of this operation, paired with evidence of whether the row failed. + fn apply(lhs: T, rhs: T) -> (T, Self::Fail); } impl CheckedPrimitiveOp for CheckedAdd { const ERROR: &'static str = "integer overflow in checked add"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.add_value(rhs), lhs.add_error(rhs)) } @@ -76,9 +55,9 @@ impl CheckedPrimitiveOp for CheckedAdd { impl CheckedPrimitiveOp for CheckedSub { const ERROR: &'static str = "integer overflow in checked sub"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.sub_value(rhs), lhs.sub_error(rhs)) } @@ -87,9 +66,9 @@ impl CheckedPrimitiveOp for CheckedSub { impl CheckedPrimitiveOp for CheckedMul { const ERROR: &'static str = "integer overflow in checked mul"; - type Failure = T::MulFailure; + type Fail = T::MulFailure; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, T::MulFailure) { (lhs.mul_value(rhs), lhs.mul_failure(rhs)) } @@ -97,16 +76,10 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { let failed = lhs.div_error(rhs); let value = if failed { @@ -116,151 +89,16 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } } -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) -} - -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. +/// Per-width arithmetic used to compute values and failure evidence. /// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// The add, subtract, and multiply value methods **must** be total over every stored lane value. +/// [`Self::div_value`] may assume that [`Self::div_error`] returned `false` for the same operands. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// This may be a word rather than `bool` when narrowing evidence would block vectorization. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -269,18 +107,15 @@ trait CheckedArithmetic: NativePType { fn sub_error(self, rhs: Self) -> bool; fn mul_value(self, rhs: Self) -> Self; fn mul_failure(self, rhs: Self) -> Self::MulFailure; + + /// Divide operands that [`Self::div_error`] accepted. fn div_value(self, rhs: Self) -> Self; + + /// Return whether [`Self::div_value`] would trap for these operands. fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } -/// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// Generate the shared integer operations from their failure predicates. macro_rules! impl_checked_integer { ( $ty:ty, @@ -291,67 +126,57 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self.wrapping_add(rhs) } - #[inline(always)] + #[inline] fn add_error(self, rhs: Self) -> bool { let ($add_lhs, $add_rhs) = (self, rhs); $add_error } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self.wrapping_sub(rhs) } - #[inline(always)] + #[inline] fn sub_error(self, rhs: Self) -> bool { let ($sub_lhs, $sub_rhs) = (self, rhs); $sub_error } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self.wrapping_mul(rhs) } - #[inline(always)] + #[inline] $(#[$mul_failure_attr])* fn mul_failure(self, rhs: Self) -> $mul_failure_ty { let ($mf_lhs, $mf_rhs) = (self, rhs); $mul_failure } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, rhs: Self) -> bool { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// Unsigned multiplication reports its discarded high half as failure evidence. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -364,12 +189,7 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// Signed widths use a range check or discarded high-half evidence. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -377,9 +197,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -389,13 +206,17 @@ macro_rules! impl_checked_signed { let kept = wide as $ty; let discarded = (wide >> <$ty>::BITS) as $ty; + // A product fits exactly when its discarded half is the sign extension of the kept + // half. XOR reduces that comparison to zero evidence for success and nonzero evidence + // for overflow without converting the wide product to a branch. + (discarded ^ (kept >> (<$ty>::BITS - 1))) as $failure }); }; ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -408,7 +229,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -417,54 +238,47 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self + rhs } - #[inline(always)] + #[inline] fn add_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self - rhs } - #[inline(always)] + #[inline] fn sub_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self * rhs } - #[inline(always)] + #[inline] fn mul_failure(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -484,30 +298,27 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; - /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change - /// pivots, and a spread of magnitudes that straddles the 64-bit split. + /// Values around zero, signed extrema, and 32- and 64-bit boundaries where the discarded + /// multiplication half or its sign extension changes. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, // Additive identity. + 1, // Smallest positive value. + -1, // All sign bits set. + 2, // Small positive power of two. + -2, // Small negative power of two. + 3, // Small non-power of two. + i64::MIN, // Minimum signed value. + i64::MIN + 1, // Minimum signed value's neighbor. + i64::MAX, // Maximum signed value. + i64::MAX - 1, // Maximum signed value's neighbor. + 1 << 31, // First positive value outside i32. + 1 << 32, // First value with bit 32 set. + 1 << 62, // Largest positive power of two in i64. + -(1 << 62), // Negative counterpart of the largest power of two. + 0x7FFF_FFFF, // Maximum i32 represented as i64. + -0x8000_0000, // Minimum i32 represented as i64. ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { let failed = lhs.mul_failure(rhs) != ::default(); @@ -522,14 +333,11 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..b6f565d8ff4 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive arithmetic execution through [`RowFn`]. +//! +//! `Binary` keeps its registered contract; [`NumericBinary`] is only an execution helper. Decimal +//! arithmetic remains on its existing columnar path. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedArithmetic; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::unstable::row::InitializedElement; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::UninitElementSink; +use crate::scalar_fn::unstable::row::execute_rows; + +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + execute_rows(&NumericBinary, &op, &args, ctx) +} + +/// Internal row execution for the primitive arithmetic operators. +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + // Fallibility is queried without input dtypes, so this conservatively covers integer widths. + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + // `NumericBinary` is a private implementation detail of `Binary`: it is never registered or + // serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. + ScalarFnVTable::id(&Binary) + } + + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = PType::try_from( + args.first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?, + )?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_div::(visitor), + } + }) + } +} + +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + visitor.visit_deferred::<(T, T), T, Op::Fail>( + |(lhs, rhs)| Op::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(numeric_error(Op::ERROR)); + } + + Ok(()) + }, + ) +} + +fn visit_div(visitor: V) -> VortexResult +where + T: CheckedArithmetic, + V: RowVisitor, +{ + if T::PTYPE.is_float() { + return visit_checked::(visitor); + } + + // Integer division is scalar and expensive, so deferring its cheap failure check preserves no + // vectorization. Check each divide immediately and stop at the first failure. + // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when + // they need to skip invalid rows. + visitor.visit_into::<(T, T), UninitElementSink, _>(|(lhs, rhs), output| { + let (value, failed) = CheckedDiv::apply(lhs, rhs); + if failed { + return Err(numeric_error(>::ERROR)); + } + + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }) +} + +/// Keep rich error construction out of row closures so the closures remain inlineable. +#[cold] +#[inline(never)] +fn numeric_error(message: &'static str) -> VortexError { + vortex_err!(InvalidArgument: "{message}") +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..3813c8612b3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -201,8 +201,7 @@ fn test_integer_array_array_errors_on_valid_lanes() { assert!(result.is_err()); } -/// Multiply two non-nullable lanes of `lhs` by two of `rhs`, expecting `Some(product)` where the -/// product fits and `None` where the checked kernel must report overflow. +/// Assert one checked multiplication through the complete array execution path. #[track_caller] fn assert_multiply(lhs: T, rhs: T, expected: Option) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -297,13 +296,11 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// An overflow late in the batch must still be reported, unless its row is null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs new file mode 100644 index 00000000000..922e3ee5095 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution arguments paired with the metadata selected during planning. +//! +//! [`BorrowedExecutionArgs`] can point at original, sliced, or filtered arrays while retaining the +//! dtypes, output dtype, and null policy of the original batch plan. + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::visitor::RowPolicy; + +/// A borrowed [`ExecutionArgs`] view with the planning metadata selected for its row kernel. +/// +/// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the +/// original planned batch. Keeping them together prevents an execution path from pairing an input +/// view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub(crate) struct BorrowedExecutionArgs<'a> { + /// The input arrays for this kernel invocation. + arrays: &'a [ArrayRef], + + /// The number of rows in this kernel invocation. + row_count: usize, + + /// The original input dtypes used to select the row implementation. + dtypes: &'a [DType], + + /// The non-nullable dtype built by the selected output capability. + output_dtype: &'a DType, + + /// The nullable execution policy selected during planning. + policy: RowPolicy, +} + +impl<'a> BorrowedExecutionArgs<'a> { + /// Pair one input view with the planning metadata selected for its batch. + pub(crate) fn new( + arrays: &'a [ArrayRef], + row_count: usize, + dtypes: &'a [DType], + output_dtype: &'a DType, + policy: RowPolicy, + ) -> Self { + Self { + arrays, + row_count, + dtypes, + output_dtype, + policy, + } + } + + /// Return the concrete arrays used by encoding-aware execution. + pub(crate) fn arrays(&self) -> &'a [ArrayRef] { + self.arrays + } + + /// Return the original input dtypes used to select the row implementation. + pub(crate) fn dtypes(&self) -> &'a [DType] { + self.dtypes + } + + /// Return the non-nullable dtype built by the selected output capability. + pub(crate) fn output_dtype(&self) -> &'a DType { + self.output_dtype + } + + /// Return the nullable execution policy selected during planning. + pub(crate) fn policy(&self) -> RowPolicy { + self.policy + } +} + +impl ExecutionArgs for BorrowedExecutionArgs<'_> { + fn get(&self, index: usize) -> VortexResult { + self.arrays.get(index).cloned().ok_or_else(|| { + vortex_err!( + "row-function input index must be less than {}, got {index}", + self.arrays.len(), + ) + }) + } + + fn num_inputs(&self) -> usize { + self.arrays.len() + } + + fn row_count(&self) -> usize { + self.row_count + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs new file mode 100644 index 00000000000..8cb89a14bf3 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::scalar_fn::unstable::row::execute::RowExecution; + +impl Batch { + /// Evaluate one row of constant inputs and broadcast the validated result. + pub(super) fn broadcast_one_row( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let result = VortexResult::from(kernel(self.execution_args(&one_row, 1), ctx)?)?; + let result = self.validate_kernel_output(result, 1, ctx)?; + let result = self.finalize_output(result, 1)?; + let scalar = result.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.row_count).into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs new file mode 100644 index 00000000000..d3e0ef103c2 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +impl Batch { + /// Run every stored payload, then attach the input validity without materializing its mask. + pub(super) fn execute_dense( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let values = match kernel(self.execution_args(&self.inputs, self.row_count), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Unlike `resolve_validity`, all-true preserves the deferred error and all-false + // suppresses evidence that came entirely from null rows. An empty loop cannot + // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Deferred retry receives only the dense kernel. Filter first so the retry cannot + // evaluate null rows again. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled by the guard in `Batch::execute`, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs new file mode 100644 index 00000000000..871c7fc1219 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +impl Batch { + /// Filter to valid rows, run the kernel, then scatter into a null-padded output. + pub(super) fn filter_and_scatter( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let values = VortexResult::from(kernel( + self.execution_args(&filtered, valid.true_count()), + ctx, + )?)?; + let values = self.validate_kernel_output(values, valid.true_count(), ctx)?; + + self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) + } + + /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set + /// bits, producing an array of length `valid.len()` that is null at every unset position. + fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { + vortex_ensure_eq!( + values.len(), + valid.true_count(), + "the {} kernel output must contain {} filtered rows, got {}", + self.id, + valid.true_count(), + values.len(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!( + "scatter_valid requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + + // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index + // 0, and any in-bounds index would do since they are masked out below (values is non-empty + // here). + let mut indices = vec![0u64; valid.len()]; + let mut rank = 0u64; + for &(start, end) in slices { + for index in &mut indices[start..end] { + *index = rank; + rank += 1; + } + } + let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); + + let scattered = values.take(indices)?; + + // A nullable gathered array cannot be wrapped because a `Masked` child must be all valid. + // The general masking pass unions its nulls with the batch validity instead. + if scattered.dtype().is_nullable() { + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + return scattered.mask(mask); + } + + // The gathered values are all valid, so attaching validity is sufficient. + Ok(MaskedArray::try_new( + scattered, + Validity::from_mask(valid.clone(), Nullability::Nullable), + )? + .into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs new file mode 100644 index 00000000000..f2d4f605740 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Selects a batch execution strategy. +//! +//! [`Batch::execute`] handles universal fast paths and encoded reductions, then delegates to dense +//! or valid-only execution. + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::Batch; +use super::RowPolicy; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Constant; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::types::batch_constant; +use crate::validity::Validity; + +mod constant; +mod dense; +mod filter_scatter; +mod valid_only; + +mod output; +pub(crate) use output::finalize_kernel_output; + +impl Batch { + /// Apply encoded reductions, constant folding, and null handling around `kernel`. + /// + /// `reduce` receives the original inputs before constant broadcasting. When the mask contains + /// valid and invalid rows, `try_unfiltered` may avoid filtering. `Ok(None)` filters the valid + /// rows and scatters the output back. Every result is checked against the planned shape and + /// dtype. + pub(crate) fn execute( + &self, + reduce: impl FnOnce( + BorrowedExecutionArgs<'_>, + &mut ExecutionCtx, + ) -> VortexResult>, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.validity, Validity::AllInvalid) + || self.inputs.iter().any(|input| { + input + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()) + }) + { + return Ok(self.all_null()); + } + + // An empty mask is both all-true and all-false, so deferred encoded evidence cannot be + // attributed to an observable row. Let the ordinary policy construct the typed empty + // output instead. + if self.row_count > 0 + && let Some(execution) = reduce(self.execution_args(&self.inputs, self.row_count), ctx)? + { + match execution { + RowExecution::Output(values) => return self.finalize_reduced(values, ctx), + RowExecution::DeferredError(error) => { + return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx); + } + } + } + + // All inputs constant, and their conjoined validity proves every row non-null. This sees + // through extension and masked wrappers just like argument decoding does. + if self.row_count > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs new file mode 100644 index 00000000000..f4dfc44fae2 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +use super::super::Batch; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + +impl Batch { + pub(super) fn all_null(&self) -> ArrayRef { + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + } + + pub(super) fn finalize_output( + &self, + values: ArrayRef, + expected_len: usize, + ) -> VortexResult { + reconcile_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Reconcile an encoding-aware result and apply the batch's strict input validity. + pub(super) fn finalize_reduced( + &self, + values: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + validate_output(self.id, &self.result_dtype, self.row_count, &values)?; + + let input_valid = self.validity.execute_mask(self.row_count, ctx)?; + let output_valid = values.validity()?.execute_mask(self.row_count, ctx)?; + vortex_ensure!( + input_valid.bitand_not(&output_valid).all_false(), + "the {} encoded reduction produced nulls for valid rows", + self.id, + ); + + let values = match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => values, + Validity::Array(valid) => values.mask(valid)?, + // Handled before the encoding-aware hook runs. + Validity::AllInvalid => return Ok(self.all_null()), + }; + + cast_output_nullability(&self.result_dtype, values) + } + + /// Validate the output from a row kernel before batch validity is attached. + pub(super) fn validate_kernel_output( + &self, + values: ArrayRef, + expected_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) + } +} + +/// Validate the output produced directly by a row kernel. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring +/// nullability, and every produced row **must** be valid. Batch execution owns strict null +/// propagation and attaches input-derived validity only after this boundary. +pub(crate) fn finalize_kernel_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + vortex_ensure!( + values.all_valid(ctx)?, + "the {id} row kernel must produce only valid rows, got at least one null row", + ); + + cast_output_nullability(result_dtype, values) +} + +/// Reconcile an output with the function's declared shape and nullability. +fn reconcile_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + + cast_output_nullability(result_dtype, values) +} + +/// Validate an output's shape and logical dtype without executing a nullability cast. +fn validate_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: &ArrayRef, +) -> VortexResult<()> { + vortex_ensure_eq!( + values.len(), + expected_len, + "the {id} kernel output must contain {expected_len} rows, got {}", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel output dtype must match {result_dtype} ignoring nullability, got {}", + values.dtype(), + ); + + Ok(()) +} + +/// Cast only the output nullability after its shape, dtype, and validity are accepted. +fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult { + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs new file mode 100644 index 00000000000..d68c2e57da8 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +/// The result of resolving batch validity. +enum ResolvedValidity { + /// The output for an all-valid or all-null batch. + Output(ArrayRef), + + /// A mask with both valid and invalid rows. + PartiallyValid(Mask), +} + +impl Batch { + /// Resolve deferred evidence from the encoded path by executing only observable rows. + pub(super) fn resolve_reduced_error( + &self, + error: VortexError, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Resolve validity, try unfiltered execution, then fall back to filtering. + pub(super) fn execute_valid_only( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedValidity::Output(output) => return Ok(output), + ResolvedValidity::PartiallyValid(valid) => valid, + }; + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Materialize validity and handle all-valid or all-null batches. + fn resolve_validity( + &self, + kernel: &impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + let values = VortexResult::from(kernel( + self.execution_args(&self.inputs, self.row_count), + ctx, + )?)?; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.finalize_output(values, self.row_count)?; + + return Ok(ResolvedValidity::Output(values)); + } + + if valid.all_false() { + return Ok(ResolvedValidity::Output(self.all_null())); + } + + Ok(ResolvedValidity::PartiallyValid(valid)) + } + + /// Try execution against the original inputs, then mask a returned full-length result. + fn try_execute_unfiltered( + &self, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(execution) = try_unfiltered( + self.execution_args(&self.inputs, self.row_count), + valid, + ctx, + )? + else { + return Ok(None); + }; + let values = VortexResult::from(execution)?; + let values = self.validate_kernel_output(values, valid.len(), ctx)?; + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + self.finalize_output(values.mask(mask)?, valid.len()) + .map(Some) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs new file mode 100644 index 00000000000..a246ba3e30b --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution around a strict row kernel. +//! +//! A row kernel handles typed values for one row. This module adds the columnar concerns around it: +//! planning the output and null strategy, preserving batch constants, propagating strict validity, +//! selecting an execution strategy, and validating the finished output. +//! +//! [`BatchPlan`] carries the nullable execution strategy selected by a concrete dispatch. [`Batch`] +//! applies that strategy, and [`BorrowedExecutionArgs`] pairs each kernel invocation with its +//! planning metadata. + +use smallvec::SmallVec; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + +mod args; +pub(super) use args::BorrowedExecutionArgs; + +mod execute; +pub(super) use execute::finalize_kernel_output; + +mod planning; + +pub(super) use super::visitor::BatchPlan; +pub(super) use super::visitor::RowPolicy; + +/// One batch of inputs and the metadata needed before its row kernel runs. +pub(crate) struct Batch { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The number of rows in the original execution scope. + row_count: usize, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs new file mode 100644 index 00000000000..24bb31b2709 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use super::Batch; +use super::BatchPlan; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + +impl Batch { + /// Collect the inputs and derive their dtypes, validity, and execution policy. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check would vacuously pass. + pub(crate) fn new( + id: ScalarFnId, + args: &dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let row_count = args.row_count(); + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + for (index, input) in inputs.iter().enumerate() { + vortex_ensure_eq!( + input.len(), + row_count, + "the {id} input {index} must have {row_count} rows, got {}", + input.len(), + ); + } + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let result_dtype = plan.result_dtype(&arg_dtypes); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + row_count, + inputs, + arg_dtypes, + validity, + result_dtype, + output_dtype: plan.output_dtype, + policy: plan.policy, + }) + } + + /// Pair an input view with this batch's planning metadata. + pub(super) fn execution_args<'b>( + &'b self, + arrays: &'b [ArrayRef], + row_count: usize, + ) -> BorrowedExecutionArgs<'b> { + BorrowedExecutionArgs::new( + arrays, + row_count, + &self.arg_dtypes, + &self.output_dtype, + self.policy, + ) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs new file mode 100644 index 00000000000..ebf69580c54 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -0,0 +1,974 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::super::execute::RowExecution; +use super::Batch; +use super::BatchPlan; +use super::RowPolicy; +use super::finalize_kernel_output; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; +use crate::scalar_fn::unstable::row::row_fn_return_dtype; +use crate::validity::Validity; + +#[derive(Clone)] +struct RetryConstantAdd; + +#[derive(Clone)] +struct NullarySeven; + +#[derive(Clone)] +struct AddThree; + +#[derive(Clone)] +struct AddShort(ShortVisit); + +/// An element whose decode drops the last row, standing in for an invalid element implementation. +struct ShortDecodeI64; + +#[derive(Clone)] +struct OriginalInputReducer; + +#[derive(Clone)] +struct InvalidEncodedReduction; + +#[derive(Clone)] +struct DeferredOriginalReducer; + +#[derive(Clone)] +struct ValidOnlyIdentity; + +#[derive(Clone)] +struct SinkOptions; + +struct OptionsCheckingSink; + +#[derive(Clone)] +struct InvalidKernelOutput; + +/// Deliberately violates [`OutputElement::build`] to test validation at the public boundary. +struct NullProducingI64(i64); + +#[derive(Clone)] +struct PreparedAdd { + visit: PreparedVisit, + prepares: Arc, +} + +#[derive(Clone, Copy)] +enum PreparedVisit { + Owned, + Sink, + Deferred, +} + +#[derive(Clone, Copy)] +enum ShortVisit { + Owned, + Sink, +} + +// SAFETY: the view and unchecked access delegate to the `i64` implementation. The implementation +// deliberately returns a short column so the executor's pre-loop length guard can be tested. +unsafe impl InputElement for ShortDecodeI64 { + type Column = Buffer; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = false; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let column = ::decode(array, ctx)?; + + Ok(column.slice(0..column.len().saturating_sub(1))) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + ::view(column) + } + + fn view_len(view: &Self::View<'_>) -> usize { + ::view_len(view) + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_from_view(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { ::get_from_view_unchecked(view, index) } + } +} + +// SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or +// `finish` through the executor. The row-initialization requirements are therefore vacuous. +unsafe impl OutputSink for OptionsCheckingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn output_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { + if !enabled { + vortex_bail!(InvalidArgument: "the test sink is disabled"); + } + + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize) -> VortexResult { + vortex_bail!("the planning-only test sink must not be allocated") + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + unsafe fn finish(self) -> VortexResult { + vortex_bail!("the planning-only test sink must not finish") + } +} + +impl OutputElement for NullProducingI64 { + fn element_dtype() -> DType { + DType::from(i64::PTYPE) + } + + fn build(values: Vec) -> ArrayRef { + let values: Vec<_> = values.into_iter().map(|value| value.0).collect(); + let validity = Validity::from_iter((0..values.len()).map(|index| index != 0)); + + PrimitiveArray::new(values, validity).into_array() + } +} + +struct I64Sink(BufferMut); + +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +impl RowFn for NullarySeven { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_seven"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), I64Sink, _>(|(), output| { + *output = 7; + }) + } +} + +impl RowFn for AddThree { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["first", "second", "third"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_three"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64, i64), i64>(|(first, second, third)| first + second + third) + } +} + +impl RowFn for AddShort { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_short"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + match self.0 { + ShortVisit::Owned => { + visitor.visit::<(ShortDecodeI64, i64), i64>(|(lhs, rhs)| lhs + rhs) + } + ShortVisit::Sink => { + visitor.visit_into::<(ShortDecodeI64, i64), I64Sink, _>(|(lhs, rhs), output| { + *output = lhs + rhs; + }) + } + } + } +} + +impl RowFn for RetryConstantAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.retry_constant_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(u8, u8), u8, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "checked add overflowed")); + } + + Ok(()) + }, + ) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 1 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(0u8, args[0].len()).into_array(), + ))); + } + + Ok(None) + } +} + +impl RowFn for OriginalInputReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.original_input_reducer"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 3 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(42_i64, 3).into_array(), + ))); + } + + Ok(None) + } +} + +impl RowFn for InvalidEncodedReduction { + type Options = usize; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_encoded_reduction"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + null_index: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::Output( + PrimitiveArray::new( + vec![10_i64, 20], + Validity::from_iter((0..2).map(|index| index != *null_index)), + ) + .into_array(), + ))) + } +} + +impl RowFn for DeferredOriginalReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.deferred_original_reducer"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::DeferredError(vortex_err!( + InvalidArgument: "encoded payload failed" + )))) + } +} + +impl RowFn for ValidOnlyIdentity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.valid_only_identity"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| { + *output = value; + Ok(()) + }) + } +} + +impl RowFn for SinkOptions { + type Options = bool; + + const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.sink_options"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), OptionsCheckingSink, _>(|(), ()| ()) + } +} + +impl RowFn for InvalidKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_kernel_output"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), NullProducingI64>(|(value,)| NullProducingI64(value)) + } +} + +impl RowFn for PreparedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.prepared_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let prepares = Arc::clone(&self.prepares); + let prepare = move |(_lhs, rhs): (Option, Option)| { + prepares.fetch_add(1, Ordering::Relaxed); + rhs + }; + + match self.visit { + PreparedVisit::Owned => visitor + .visit_prepared::<(i64, i64), i64, _>(prepare, |constant_rhs, (lhs, rhs)| { + lhs.wrapping_add(constant_rhs.unwrap_or(rhs)) + }), + PreparedVisit::Sink => visitor.visit_prepared_into::<(i64, i64), I64Sink, _, ()>( + prepare, + |constant_rhs, (lhs, rhs), output| { + *output = lhs.wrapping_add(constant_rhs.unwrap_or(rhs)); + }, + ), + PreparedVisit::Deferred => visitor.visit_prepared_deferred::<(i64, i64), i64, _, bool>( + prepare, + |constant_rhs, (lhs, rhs)| lhs.overflowing_add(constant_rhs.unwrap_or(rhs)), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "prepared add overflowed")); + } + + Ok(()) + }, + ), + } + } +} + +#[test] +fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_batch"); + + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let result = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::Dense, + }) + }); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_short_decode_beside_constant_is_rejected() -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter(0..64_i64).into_array(); + let rhs = ConstantArray::new(10_i64, 64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 64); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(ShortVisit::Sink), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short decoded column passed the pre-loop length check"), + }; + + assert!( + error + .to_string() + .contains("does not address exactly 64 rows"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[rstest] +#[case::owned(ShortVisit::Owned)] +#[case::sink(ShortVisit::Sink)] +fn test_short_constant_decode_is_rejected(#[case] visit: ShortVisit) -> VortexResult<()> { + let lhs = ConstantArray::new(10_i64, 64).into_array(); + let rhs = PrimitiveArray::from_iter(0..64_i64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 64); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(visit), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short constant decode passed the pre-loop length check"), + }; + + assert!( + error + .to_string() + .contains("batch-constant input must contain exactly 1 row, got 0"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[test] +fn test_short_constant_null_tolerant_decode_is_rejected() -> VortexResult<()> { + let lhs = MaskedArray::try_new( + ConstantArray::new(10_i64, 4).into_array(), + Validity::from_iter([true, false, true, false]), + )? + .into_array(); + let rhs = PrimitiveArray::from_iter(0..4_i64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 4); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(ShortVisit::Sink), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short null-tolerant constant decode passed validation"), + }; + + assert!( + error + .to_string() + .contains("decoded batch-constant input must contain exactly 1 row, got 0"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[test] +fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("valid-row overflow must remain observable"), + }; + + assert!( + error.to_string().contains("checked add overflowed"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[test] +fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![1, u8::MAX], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1_u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_u8, 0], Validity::from_iter([true, false])); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_defers_errors_behind_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_empty_batch_skips_deferred_encoded_error() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(Vec::::new()).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid(Validity::AllValid)] +#[case::mixed(Validity::from_iter([true, false]))] +fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![10_i64, 20], validity).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&InvalidEncodedReduction, &0, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an encoded reduction introduced a null on a valid row"), + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_encoded_reduction"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("encoded reduction produced nulls for valid rows"), + "the boundary error must identify invalid reduced output, got {error}", + ); + Ok(()) +} + +#[test] +fn test_reduce_encoded_preserves_input_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&InvalidEncodedReduction, &1, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let expected = ConstantArray::new(42_i64, 3).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid([true, true])] +#[case::all_invalid([false, false])] +fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.resolve_validity"); + + let validity = Validity::Array(BoolArray::from_iter(validity).into_array()); + let input = PrimitiveArray::new(vec![4_i64, 5], validity).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_filters_and_scatters() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.filter_and_scatter"); + + let input = PrimitiveArray::new( + vec![10_i64, 20, 30, 40], + Validity::from_iter([true, false, true, false]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 4); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.finalize_kernel_output"); + + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let result_dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + let mut ctx = array_session().create_execution_ctx(); + + let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone(), &mut ctx)?; + let expected = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + assert_eq!(actual.dtype(), &result_dtype); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + assert!(finalize_kernel_output(*ID, &result_dtype, 3, values, &mut ctx).is_err()); + + let bools = BoolArray::from_iter([true, false]).into_array(); + assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools, &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn test_output_dtype_receives_function_options() -> VortexResult<()> { + assert_eq!( + row_fn_return_dtype(&SinkOptions, &true, &[])?, + DType::from(i64::PTYPE) + ); + assert!(row_fn_return_dtype(&SinkOptions, &false, &[]).is_err()); + Ok(()) +} + +#[rstest] +#[case::nonnullable(Validity::NonNullable)] +#[case::all_valid(Validity::AllValid)] +fn test_kernel_output_rejects_nulls_at_function_boundary( + #[case] validity: Validity, +) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + let execution = execute_rows(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); + let error = match execution { + Err(error) => error, + Ok(output) => match output.execute::(&mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an invalid row kernel output passed boundary validation"), + }, + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_kernel_output"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("row kernel must produce only valid rows"), + "the boundary error must identify invalid row output, got {error}", + ); + Ok(()) +} + +#[rstest] +#[case::owned_constant(PreparedVisit::Owned, true)] +#[case::owned_per_row(PreparedVisit::Owned, false)] +#[case::sink_constant(PreparedVisit::Sink, true)] +#[case::sink_per_row(PreparedVisit::Sink, false)] +#[case::deferred_constant(PreparedVisit::Deferred, true)] +#[case::deferred_per_row(PreparedVisit::Deferred, false)] +fn test_prepared_visits( + #[case] visit: PreparedVisit, + #[case] constant_rhs: bool, +) -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let rhs = if constant_rhs { + ConstantArray::new(3_i64, 2).into_array() + } else { + PrimitiveArray::from_iter([3_i64, 4]).into_array() + }; + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let prepares = Arc::new(AtomicUsize::new(0)); + let function = PreparedAdd { + visit, + prepares: Arc::clone(&prepares), + }; + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = if constant_rhs { + PrimitiveArray::from_iter([4_i64, 5]).into_array() + } else { + PrimitiveArray::from_iter([4_i64, 6]).into_array() + }; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(prepares.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn test_nullary_row_function_broadcasts() -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([7_i64, 7, 7]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_empty_batch_preserves_nonnullable_dtype() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(std::iter::empty::()).into_array(); + let args = VecExecutionArgs::new(vec![input], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&ValidOnlyIdentity, &EmptyOptions, &args, &mut ctx)?; + + assert_eq!(actual.len(), 0); + assert_eq!(actual.dtype(), &DType::from(i64::PTYPE)); + Ok(()) +} + +#[test] +fn test_owned_execution_traverses_three_per_row_inputs() -> VortexResult<()> { + let args = VecExecutionArgs::new( + vec![ + PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(), + PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(), + PrimitiveArray::from_iter([100_i64, 200, 300]).into_array(), + ], + 3, + ); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&AddThree, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([111_i64, 222, 333]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs new file mode 100644 index 00000000000..73722549f41 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-loop execution for owned outputs and output sinks. +//! +//! [`owned`] stores one independent value per row and reduces compact failure evidence. [`sink`] +//! drives output builders whose row handles may share batch state. Both return [`RowExecution`], +//! which distinguishes a completed array from a deferred error that batch validity may suppress. + +mod owned; +pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_infallible; + +mod outcome; +pub use outcome::RowExecution; + +mod sink; +pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_valid_rows; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs new file mode 100644 index 00000000000..fc013e7a317 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The result of a completed row loop before batch-level null handling. +//! +//! [`RowExecution`] preserves deferred failure evidence until batch execution can determine whether +//! the failing payload belonged to a valid row. + +use vortex_error::VortexError; +use vortex_error::VortexResult; + +use crate::ArrayRef; + +/// The outcome of a row loop before batch execution decides whether an error is observable. +/// +/// Together with the surrounding [`VortexResult`], this represents three outcomes: +/// +/// - `Err(error)` is a non-retryable execution or immediate row error. +/// - [`Output`](Self::Output) is a successful row loop. +/// - [`DeferredError`](Self::DeferredError) is failure evidence from a completed row loop. +/// +/// A dense loop can evaluate null payloads, so its deferred error is not always observable. Batch +/// execution can retry only valid rows to discard errors caused by null payloads. A plain +/// `VortexResult` cannot distinguish these errors from failures that a retry cannot fix. +/// +/// Once execution is known to contain only valid rows, converting this outcome into a +/// `VortexResult` turns [`DeferredError`](Self::DeferredError) into an ordinary error. +pub enum RowExecution { + /// The successfully built, full-length output column. + Output(ArrayRef), + + /// An error constructed from failure evidence reduced across a completed row loop. + DeferredError(VortexError), +} + +impl From for VortexResult { + fn from(execution: RowExecution) -> Self { + match execution { + RowExecution::Output(output) => Ok(output), + RowExecution::DeferredError(error) => Err(error), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs new file mode 100644 index 00000000000..d3809f939be --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -0,0 +1,121 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Executes row kernels that return one independent owned value per row. +//! +//! [`execute_owned`] decodes inputs once, prepares constant state, writes into spare vector +//! capacity, and reduces compact failure evidence without putting error construction in the hot +//! loop. [`execute_owned_infallible`] removes that failure path for infallible kernels. + +use std::ops::BitOrAssign; + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; + +/// Zero-sized failure accumulator for infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column, then store one output per row from an infallible kernel. +pub(crate) fn execute_owned_infallible( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned::( + args, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode every input column, then store outputs and combine per-row failure evidence. +pub(crate) fn execute_owned( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, +{ + // The output vector stays at length zero until every slot is initialized so that an unwind + // abandons partially initialized spare capacity. This no-drop assertion proves that no + // initialized value requires a destructor to run. + const { assert_owned_output_needs_no_drop::() }; + + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let output = &mut values.spare_capacity_mut()[..row_count]; + + let failure = if let Some(views) = Args::views_no_constants(&columns) { + // Keep this validation beside the views so LLVM sees their common length here. + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + let source = unsafe { Args::indexed_source(views, row_count) }; + + source.map_checked_into(output, |elements| apply(&prepared, elements)) + } else { + // Keep this proof branch-local. Shared validation prevents LLVM from specializing this + // loop for each batch-constant arrangement, leaving it scalar under multiple CGUs without + // LTO. The exact pass interaction is unknown. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated = Fail::default(); + + // Iterate over `output` directly. A `0..row_count` range reuses the address-taken value + // from the validation error formatter and retains an output bounds check. + for (index, slot) in output.iter_mut().enumerate() { + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the loop. + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + + slot.write(value); + accumulated |= row_failure; + } + + accumulated + }; + + // SAFETY: normal completion of either execution path initializes `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Defer failures so batch execution can retry with only valid rows. + match finish_failure(failure) { + Ok(()) => Ok(RowExecution::Output(Out::build(values))), + Err(error) => Ok(RowExecution::DeferredError(error)), + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs new file mode 100644 index 00000000000..4fd2b42476b --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -0,0 +1,441 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Executes row kernels that write through an [`OutputSink`]. +//! +//! Dense execution visits every row. Skip-invalid execution initializes skipped output rows and +//! visits only rows that are valid in every input. Skip-invalid execution declines when either the +//! input representation or sink cannot support that path. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::SinkResult; + +/// Verify that every decoded input addresses exactly `row_count` rows. +/// +/// Unlike the owned executor, the paths with and without batch constants can share this check +/// without losing sink-loop vectorization under multiple CGUs without LTO. The exact pass +/// interaction is unknown. +fn verify_lengths( + columns: &Args::Columns, + views: Option<&Args::Views<'_>>, + row_count: usize, +) -> VortexResult<()> { + let lengths_match = match views { + Some(views) => Args::view_lens_match(views, row_count), + None => Args::decoded_lens_match(columns, row_count), + }; + + vortex_ensure!( + lengths_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(()) +} + +/// Decode inputs once, then write one sink row for each input row. +/// +/// The executor owns the sink and passes each output row to `apply`. This keeps `apply` as [`Fn`]. +/// Capturing the sink would require [`FnMut`] and put its buffer metadata behind loop-carried +/// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant. +pub(crate) fn execute_sink( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + let columns = Args::decode(args, ctx)?; + let views = Args::views_no_constants(&columns); + + let row_count = args.row_count(); + verify_lengths::(&columns, views.as_ref(), row_count)?; + + let constants = Args::constants(&columns); + let prepared = prepare(constants); + + let mut sink = >::with_capacity(row_count)?; + + // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. + { + let mut rows = >::rows(&mut sink); + + // This equality proves to LLVM that `0..row_count` is in bounds for `rows`. + let sink_row_count = >::row_count(&rows); + vortex_ensure_eq!( + sink_row_count, + row_count, + "the output sink must address exactly {row_count} rows, got {sink_row_count}", + ); + + if let Some(views) = views { + for index in 0..row_count { + // SAFETY: `verify_lengths` proved every view has `row_count` rows before the loop. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + apply(&prepared, elements, output).into_result()?; + } + } else { + for index in 0..row_count { + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the + // loop. + apply(&prepared, Args::get(&columns, index), output).into_result()?; + } + } + } + + // SAFETY: every row callback completed successfully, so each returned the required write token. + unsafe { >::finish(sink) }.map(RowExecution::Output) +} + +/// Write only the rows set in `valid`, or decline when the inputs or sink cannot support +/// skip-invalid execution. +/// +/// `Ok(None)` signals batch execution to filter every input to the valid rows, run the dense +/// kernel, and scatter the results back into a null-padded array. +pub(crate) fn execute_sink_valid_rows( + args: &dyn ExecutionArgs, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult> +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + let Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + mut sink, + }) = setup_sink_valid_rows::(args, valid, ctx)? + else { + return Ok(None); + }; + + let views = Args::views_no_constants(&columns); + verify_lengths::(&columns, views.as_ref(), row_count)?; + + let constants = Args::constants(&columns); + let prepared = prepare(constants); + + // Keep `rows` scoped so its borrow ends before `finish`. With multiple CGUs and no LTO, using + // `drop(rows)` duplicates `Args::get` in every sparse callback. + { + // Initialize every slot before visiting only valid rows. + let mut rows = >::rows(&mut sink); + initialize_skipped_rows(&mut rows); + + // The initializer can change addressability. Recheck it so LLVM can prove every mask + // index is in bounds. + let initialized_row_count = >::row_count(&rows); + vortex_ensure_eq!( + initialized_row_count, + row_count, + "the initialized output sink must address exactly {row_count} rows, got {initialized_row_count}", + ); + + if let Some(views) = views { + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check proved that the sink addresses + // every mask index, which is below the mask's validated `row_count`. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + // SAFETY: `verify_lengths` proved every view has `row_count` rows, and mask indices + // are below `row_count`. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + + apply(&prepared, elements, output).into_result() + })?; + } else { + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check proved that the sink addresses + // every mask index, which is below the mask's validated `row_count`. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + apply(&prepared, Args::get(&columns, index), output).into_result() + })?; + } + } + + // SAFETY: the initializer completed before traversal, and every visited callback completed + // successfully and returned the required write token. + unsafe { >::finish(sink) } + .map(RowExecution::Output) + .map(Some) +} + +/// State resolved before preparing the skip-invalid row loop. +struct ValidRowsSetup<'valid, Args, Sink, Options> +where + Args: ElementTuple, + Sink: OutputSink, +{ + initialize_skipped_rows: for<'rows> fn(&mut >::Rows<'rows>), + columns: Args::Columns, + valid_rows: &'valid BitBuffer, + row_count: usize, + sink: Sink, +} + +/// Resolve the capabilities, inputs, sink, and validity mask for skip-invalid execution. +fn setup_sink_valid_rows<'valid, Args, Sink, Options>( + args: &dyn ExecutionArgs, + valid: &'valid Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult>> +where + Args: ElementTuple, + Sink: OutputSink, +{ + // The initializer both declares support for skipping rows and initializes those rows. + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering. Decline when any input + // cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + + let row_count = args.row_count(); + + // Keep allocation before the validity and length checks. With multiple CGUs and no LTO, + // moving it later inlines `Args::get` into every sparse callback, duplicating its bounds + // checks. + let sink = >::with_capacity(row_count)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid_rows) = valid.bit_buffer() else { + vortex_bail!( + "execute_sink_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + vortex_ensure_eq!( + valid_rows.len(), + row_count, + "the validity mask must address exactly {row_count} rows, got {}", + valid_rows.len(), + ); + + Ok(Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + sink, + })) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::RowExecution; + use super::execute_sink_valid_rows; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::NativePType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::InitializedElement; + use crate::scalar_fn::unstable::row::OutputSink; + use crate::scalar_fn::unstable::row::UninitElementSink; + use crate::validity::Validity; + + struct NonSkippingSink; + + struct ShrinkingSink(Vec); + + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or + // `finish` through the executor. The row-initialization requirements are therefore vacuous. + unsafe impl OutputSink for NonSkippingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize) -> VortexResult { + Err(vortex_err!( + "a non-skipping sink must decline before allocation" + )) + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { + } + + unsafe fn finish(self) -> VortexResult { + Err(vortex_err!("a non-skipping sink must not finish")) + } + } + + // SAFETY: the initializer deliberately shrinks the row collection to exercise the executor's + // post-initialization length check. If execution incorrectly continues, safe indexing in + // `row_unchecked` panics instead of accessing invalid memory. + unsafe impl OutputSink for ShrinkingSink { + type Rows<'a> = &'a mut Vec; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + rows.pop(); + }) + } + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(vec![0; rows])) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.0 + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::from_iter(self.0).into_array()) + } + } + + #[test] + fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([true, false]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( + &args, + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + assert!(execution.is_none()); + + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_initializes_and_writes_addressed_rows() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let valid = Mask::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::< + (i64,), + (), + UninitElementSink, + InitializedElement, + EmptyOptions, + >( + &args, + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + // SAFETY: `output` is the row supplied to this callback. + unsafe { InitializedElement::write(output, value * 2) } + }, + )?; + let Some(RowExecution::Output(actual)) = execution else { + vortex_bail!("the skip-invalid sink must produce an output"); + }; + let expected = PrimitiveArray::from_iter([20_i64, 0, 60]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_rechecks_rows_after_initialization() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( + &args, + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + *output = value; + }, + ); + + let error = match result { + Err(error) => error, + Ok(_) => vortex_bail!("the sink must reject rows changed by its initializer"), + }; + assert!( + error + .to_string() + .contains("initialized output sink must address exactly 2 rows, got 1"), + "unexpected error: {error}", + ); + + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index bcb3a008488..91f225ca071 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -11,10 +11,21 @@ //! [`RowFn::dispatch`] implementation uses a [`RowVisitor`] to select an [`ElementTuple`] and //! either an [`OutputElement`] or [`OutputSink`] for each supported dtype combination. //! +//! Unlike a general strict function, a [`RowFn`] cannot produce null from valid inputs. +//! +//! A _partially valid_ batch contains both valid and invalid rows. _Skip-invalid_ runs the kernel +//! only for valid rows without changing row positions. _Filter-and-scatter_ compacts valid rows, +//! runs the kernel, and restores their positions. +//! //! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits //! reduce compact failure evidence in that loop and retry only valid rows when null payloads may //! have caused the failure. +mod execute; +pub use execute::RowExecution; + +mod batch; + mod row_fn; pub use row_fn::RowFn; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 8a982c4fb37..93da85db841 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -4,8 +4,8 @@ //! The [`RowFn`] contract for scalar functions whose natural kernel computes one row at a time. //! //! Implementations declare their arity and fallibility, then use [`RowFn::dispatch`] to select the -//! typed row signature for each supported dtype combination. Optional methods provide -//! serialization without putting persistence plumbing in the row kernel. +//! typed row signature for each supported dtype combination. Optional hooks provide serialization +//! and encoding-aware execution without putting columnar plumbing in the row kernel. use std::fmt::Debug; use std::fmt::Display; @@ -16,10 +16,18 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use super::visitor::RowVisitor; +use crate::ArrayRef; +use crate::ExecutionCtx; use crate::dtype::DType; use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::unstable::row::RowExecution; -/// A scalar function computed one row at a time. +/// A strict scalar function whose row kernel cannot produce null from valid inputs. +/// +/// This is stronger than +/// [`ScalarFnVTable::is_strict`](crate::scalar_fn::ScalarFnVTable::is_strict), which requires null +/// propagation but permits valid inputs to produce null. The framework derives output validity +/// only from input validity. /// /// Declare argument names and use [`dispatch`](Self::dispatch) to select element and output types. /// Every implementation receives the standard [`ScalarFnVTable`]. A public type that needs custom @@ -35,12 +43,14 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// The arguments in display order. Its length is the function's exact arity. const ARG_NAMES: &'static [&'static str]; - /// Whether any dispatch can raise a semantic error. + /// Whether any dispatch or encoded reduction can raise a semantic error. /// /// See [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible) for a /// more detailed explanation of semantic errors. /// - /// The framework checks dispatched element and result types. A conservative `true` is allowed. + /// The framework checks dispatched element and result types, but cannot inspect + /// [`reduce_encoded`](Self::reduce_encoded). Set this to `true` when that hook can return a + /// semantic error or [`RowExecution::DeferredError`]. A conservative `true` is allowed. const FALLIBLE: bool; /// Returns the ID of the scalar function. @@ -71,4 +81,30 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { args: &[DType], visitor: V, ) -> VortexResult; + + /// Try an encoding-aware implementation before decoding the inputs into row elements. + /// + /// `None` continues to the row loop. [`Output`](RowExecution::Output) may remain encoded or + /// lazy. [`DeferredError`](RowExecution::DeferredError) retries only valid rows. Batch execution + /// calls this hook at most once with the original nonempty inputs. Nullary functions, empty + /// batches, slices, and compacted retries skip it. + /// + /// Like a dense row closure, this hook must be total over every stored payload, including + /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or + /// retried through the row layer. + /// + /// # Requirements + /// + /// - `output.len()` **must** equal `args[0].len()`. + /// - The output dtype **must** match the planned dtype when ignoring nullability. + /// - The output **must not** introduce a null where every input is valid. + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + _ = (options, args, ctx); + Ok(None) + } } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 2764840da7a..c9140c00b8c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -68,7 +68,7 @@ pub unsafe trait InputElement: 'static { /// cannot decode this particular array. /// /// Override this for a non-dense-safe representation that can still place safe placeholders in - /// null slots. The skip-invalid executor never reads those slots. + /// null slots. Valid-row execution never reads those slots. fn decode_null_tolerant( array: ArrayRef, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 51d66594332..7a8a8e92e41 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -20,3 +20,4 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; +pub use tuple::batch_constant; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index b0a3e709696..7c1baed168a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -27,20 +27,31 @@ pub struct ArgColumn( ); enum ArgColumnKind { + /// One decoded value per batch row; executors validate the exact length before traversal. PerRow(T::Column), + + /// Exactly one decoded row, established by [`ArgColumn::try_from_constant`]. Constant(T::Column), } impl ArgColumn { + fn try_from_constant(column: T::Column) -> VortexResult { + let decoded_len = T::view_len(&T::view(&column)); + vortex_ensure_eq!( + decoded_len, + 1, + "a decoded batch-constant input must contain exactly 1 row, got {decoded_len}", + ); + + Ok(Self(ArgColumnKind::Constant(column))) + } + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // An empty input has no row 0 to slice, and its row loop runs zero times either way. if let Some(constant) = batch_constant(&array) && !array.is_empty() { - return Ok(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?); } Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) @@ -52,10 +63,7 @@ impl ArgColumn { if let Some(constant) = batch_constant(&array) && !array.is_empty() { - return Ok(Some(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?)))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?).map(Some); } Ok(T::decode_null_tolerant(array, ctx)? @@ -88,14 +96,14 @@ impl ArgColumn { } fn addresses_rows(&self, row_count: usize) -> bool { - // A constant is always read at index zero, so it addresses any batch length. + // A constant is validated when constructed and is always read at index zero. match &self.0 { ArgColumnKind::PerRow(column) => T::view_len(&T::view(column)) == row_count, ArgColumnKind::Constant(_) => true, } } - fn constant(&self) -> Option> { + fn constant_value(&self) -> Option> { match &self.0 { ArgColumnKind::PerRow(_) => None, ArgColumnKind::Constant(column) => Some(T::get(column, 0)), @@ -130,7 +138,7 @@ pub trait ElementTuple: 'static + private::Sealed { /// The decoded column representations. type Columns; - /// Borrowed views of decoded columns when every argument stores one value per row. + /// Borrowed views of decoded columns with no batch constants. type Views<'a>; /// The borrowed row of element values. @@ -138,9 +146,9 @@ pub trait ElementTuple: 'static + private::Sealed { /// The batch-constant element values. /// - /// `Some` carries the value of a batch-constant argument. `None` marks a per-row argument. A - /// [`RowVisitor`] passes these values to its prepare closure so constant work can leave the row - /// loop. + /// `Some` carries the value of a batch-constant argument. `None` marks a non-constant argument. + /// A [`RowVisitor`] passes these values to its prepare closure so constant work can leave the + /// row loop. /// /// [`RowVisitor`]: crate::scalar_fn::unstable::row::RowVisitor type ConstElems<'a>; @@ -170,34 +178,38 @@ pub trait ElementTuple: 'static + private::Sealed { /// Decode every input column once while tolerating null rows. /// - /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid - /// strategy calls this once per batch. + /// Return `Ok(None)` when an argument has no null-tolerant representation. Valid-row execution + /// calls this once per batch. fn decode_null_tolerant( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult>; /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + /// + /// Each argument selects either its batch-constant value or row `index`. Keep that selection + /// visible in the loop so LLVM can unswitch it before vectorizing. fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; - /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// Borrow the decoded columns when none is batch-constant. /// - /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple - /// gives the optimizer ordinary contiguous column access without a per-row constant check. - fn per_row_views(columns: &Self::Columns) -> Option>; + /// Returns `None` if any column is batch-constant. Otherwise, omitting [`ArgColumn`] from the + /// returned tuple removes constant checks from the row loop. + fn views_no_constants(columns: &Self::Columns) -> Option>; /// Whether every view contains exactly `row_count` rows. /// - /// The executor calls this once before the all-per-row hot loop. A successful check gives LLVM - /// a dominating equality between the loop bound and every source length, which lets it optimize - /// the tuple access as one fixed-length traversal. + /// The executor calls this once before the loop used when no input is batch-constant. A + /// successful check gives LLVM a dominating equality between the loop bound and every source + /// length, which lets it optimize the tuple access as one fixed-length traversal. fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; - /// Whether every per-row argument contains exactly `row_count` rows. + /// Whether every non-constant argument contains exactly `row_count` rows. /// - /// This is the mixed-shape equivalent of [`view_lens_match`](Self::view_lens_match) when - /// [`per_row_views`](Self::per_row_views) declines. It runs once before the hot loop for the - /// same LLVM optimization. A batch constant is exempt because decoding collapsed it to one row. + /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include + /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch + /// constant is exempt because its [`ArgColumn`] constructor already validated the one-row + /// representation produced by decoding. fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; /// Read one row from borrowed views. @@ -257,7 +269,7 @@ impl ElementTuple for () { fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} - fn per_row_views(_columns: &Self::Columns) -> Option> { + fn views_no_constants(_columns: &Self::Columns) -> Option> { Some(()) } @@ -341,7 +353,7 @@ macro_rules! element_tuple { ($(columns.$idx.get(index),)+) } - fn per_row_views(columns: &Self::Columns) -> Option> { + fn views_no_constants(columns: &Self::Columns) -> Option> { Some(($($t::view(columns.$idx.per_row_column()?),)+)) } @@ -372,7 +384,7 @@ macro_rules! element_tuple { } fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { - ($(columns.$idx.constant(),)+) + ($(columns.$idx.constant_value(),)+) } } }; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs index d612d874935..832dbf9ee19 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -17,7 +17,7 @@ use crate::scalar_fn::unstable::row::InputElement; /// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's /// unchecked view access after batch execution validates every decoded column length once. pub trait IndexedElementTuple: ElementTuple { - /// The source shared execution uses for a dense all-per-row loop. + /// The source used when no input is batch-constant. /// /// Its length must be the common view length. For every valid index it must preserve row order, /// return the same value as [`ElementTuple::get_from_views`], and uphold the unchecked read diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index a2c143704a0..69b5cf686f6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -8,6 +8,7 @@ mod element_tuple; pub use element_tuple::ElementTuple; +pub use element_tuple::batch_constant; mod indexed; pub use indexed::IndexedElementTuple; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index 044ad15fd4b..560e1e48f4f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use super::ElementTuple; -use super::element_tuple::batch_constant; +use super::batch_constant; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index ce119f32915..e47f195410c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -12,6 +12,7 @@ pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; +pub(super) use element::batch_constant; mod result; pub use result::SinkResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index 6cc3ce06d30..22c90b8da02 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -21,8 +21,7 @@ use crate::scalar_fn::unstable::row::OutputElement; /// batch state. The executor passes each row slot into an [`Fn`] closure. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. -/// Skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an -/// initializer. +/// Execution can omit invalid rows when [`skipped_rows_initializer`] returns an initializer. /// /// # Errors /// @@ -73,10 +72,11 @@ pub unsafe trait OutputSink: 'static + Sized { /// **must not** be able to construct one without establishing the invariant. type WriteToken: 'static; - /// The operation that initializes every output position before skip-invalid execution. + /// The operation that initializes every output position before + /// [skip-invalid execution](crate::scalar_fn::unstable::row). /// - /// `Some` enables skip-invalid execution. The initializer **must** make every row safe to - /// finish. Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// `Some` enables this strategy. The initializer **must** make every row safe to finish. + /// Callbacks overwrite valid rows, and batch execution masks skipped rows. /// /// `None` makes the executor fall back to filtering the inputs. fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { @@ -153,7 +153,7 @@ impl InitializedElement { /// The row closure must return the [`InitializedElement`] from [`InitializedElement::write`] on /// success. The token is zero-sized, so the proof adds no runtime row state. /// -/// Skip-invalid execution initializes placeholders before omitting rows. Errors and unwinds are +/// When execution omits invalid rows, it initializes placeholders first. Errors and unwinds are /// safe because `values` keeps length zero until `finish`. The `T: Copy` bound means that /// initialized spare-capacity elements require no destruction. pub struct UninitElementSink { diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs new file mode 100644 index 00000000000..b003f70a395 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -0,0 +1,306 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visitors that execute dense and skip-invalid row loops. +//! +//! Each visit revalidates its concrete signature and checks that its output dtype and null policy +//! match the plan before entering a row loop. [`ExecuteValidRows`] can decline, so the batch layer +//! filters the inputs and retries with [`ExecuteRows`]. + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; + +use super::RowPolicy; +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::row_visitor::private; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_sink; +use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; + +/// The runtime visit that decodes every column once and runs the selected row loop. +pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { + /// The inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, + output_dtype: &'args DType, + policy: RowPolicy, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + options, + output_dtype, + policy, + ctx, + } + } +} + +impl private::Sealed for ExecuteRows<'_, '_, F> {} + +impl RowVisitor for ExecuteRows<'_, '_, F> { + type VisitResult = RowExecution; + + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + execute_owned_infallible::(self.args, self.ctx, prepare, apply) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink::( + self.args, self.ctx, prepare, apply, + ) + } + + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; + + execute_owned::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} + +/// The runtime visit that executes valid rows over the original input columns. +/// +/// Only output sinks have a contract for skipped output positions. Owned visits therefore decline +/// so batch execution can filter the valid inputs and scatter the output back. +pub(crate) struct ExecuteValidRows<'args, 'ctx, F: RowFn> { + /// The original inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The conjoined validity, containing both valid and invalid rows. + valid: &'args Mask, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, + output_dtype: &'args DType, + policy: RowPolicy, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + options, + output_dtype, + policy, + valid, + ctx, + } + } +} + +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} + +impl RowVisitor for ExecuteValidRows<'_, '_, F> { + type VisitResult = Option; + + fn visit_prepared( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + // Owned execution has no sink that can initialize skipped output positions. Decline so + // batch execution filters the inputs and retries with the dense visitor. + Ok(None) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink_valid_rows::( + self.args, self.valid, self.ctx, prepare, apply, + ) + } + + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + _finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; + + // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. + Ok(None) + } +} + +fn ensure_plan( + planned_output: &DType, + planned_policy: RowPolicy, + actual_output: DType, + actual_policy: RowPolicy, +) -> VortexResult<()> { + vortex_ensure_eq!( + actual_policy, + planned_policy, + "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", + ); + vortex_ensure_eq!( + actual_output, + *planned_output, + "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", + ); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index 57da5f4691b..c7f9baf6a62 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -6,9 +6,16 @@ //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch mod check; +pub(super) use check::assert_owned_output_needs_no_drop; + +mod execute; +pub(super) use execute::ExecuteRows; +pub(super) use execute::ExecuteValidRows; mod plan; +pub(super) use plan::BatchPlan; pub(super) use plan::BatchPlanner; +pub(super) use plan::RowPolicy; mod row_visitor; pub use row_visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index c522376d491..5360ced573b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -114,15 +114,12 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. - // TODO(connor)[RowFn]: Remove this allowance when the execution backend from #9130 consumes - // this policy. - #[allow(dead_code)] pub(crate) policy: RowPolicy, } impl BatchPlan { /// Return the output dtype widened with strict input nullability. - pub(crate) fn result_dtype(self, args: &[DType]) -> DType { + pub(crate) fn result_dtype(&self, args: &[DType]) -> DType { let nullability = self.output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); @@ -139,7 +136,7 @@ pub(crate) enum RowPolicy { /// Evaluate all rows, retrying only valid rows if a deferred error is raised. DenseWithRetry, - /// Execute only valid rows, trying skip-invalid execution before filtering. + /// Execute only valid rows over the original inputs before filtering. ValidOnly, } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 0e7abaa4322..1a9e8628bc7 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -74,7 +74,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// # Examples /// /// Test whether each string occurs in its allowed-values list. The prepare closure builds one - /// lookup table for a batch-constant list. The row closure scans a varying list directly. + /// lookup table for a batch-constant list. The row closure scans a per-row list directly. /// /// ```ignore /// visitor.visit_prepared::< diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index b6c14585c48..ac68ce46711 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -4,12 +4,13 @@ //! Adapts [`RowFn`] implementations to the scalar-function interface. //! //! The blanket [`ScalarFnVTable`] implementation supplies common arity, validity, fallibility, and -//! execution behavior. [`row_fn_return_dtype`] and [`execute_rows`] expose the same planning and -//! execution paths to public vtables that delegate to a private row kernel. +//! execution behavior. The visitor layer validates and executes the concrete signature selected by +//! dispatch. [`row_fn_return_dtype`] and [`execute_rows`] expose the same paths to public vtables +//! that delegate to a private row kernel. use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; use vortex_session::VortexSession; use super::row_fn::RowFn; @@ -24,6 +25,12 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::unstable::row::batch::Batch; +use crate::scalar_fn::unstable::row::batch::BorrowedExecutionArgs; +use crate::scalar_fn::unstable::row::batch::finalize_kernel_output; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::visitor::ExecuteRows; +use crate::scalar_fn::unstable::row::visitor::ExecuteValidRows; impl ScalarFnVTable for F { type Options = F::Options; @@ -69,6 +76,8 @@ impl ScalarFnVTable for F { union_child_validities(expression) } + // `RowFn` is stricter than `ScalarFnVTable::is_strict`: its kernel cannot produce null from + // valid inputs, so batch execution derives output validity only from input validity. fn is_strict(&self, _options: &Self::Options) -> bool { true } @@ -98,20 +107,40 @@ pub fn row_fn_return_dtype( /// delegate row execution to a private `RowFn` kernel through this function. pub fn execute_rows( function: &F, - _options: &F::Options, + options: &F::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { ensure_arity(function, args.num_inputs())?; - // TODO(connor)[RowFn]: Replace this temporary error with the execution backend in #9129. - vortex_bail!( - "Row function {} does not yet have an execution backend", - RowFn::id(function) + // Nullary functions have no input validity to propagate, so they skip batch execution. + if args.num_inputs() == 0 { + let plan = function.dispatch(options, &[], BatchPlanner::::new(&[], options))?; + let result_dtype = plan.result_dtype(&[]); + let nullary_args = + BorrowedExecutionArgs::new(&[], args.row_count(), &[], &plan.output_dtype, plan.policy); + + let execution = execute_row_kernel(function, options, nullary_args, ctx)?; + let values = VortexResult::from(execution)?; + + return finalize_kernel_output( + RowFn::id(function), + &result_dtype, + args.row_count(), + values, + ctx, + ); + } + + let batch = prepare_batch(function, options, args)?; + batch.execute( + |args, ctx| function.reduce_encoded(options, args.arrays(), ctx), + |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, valid, ctx| try_execute_rows_unfiltered(function, options, args, valid, ctx), + ctx, ) } -/// Validate the number of arguments before calling user-defined dispatch code. fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { let expected = F::ARG_NAMES.len(); vortex_ensure_eq!( @@ -124,26 +153,101 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +fn execute_row_kernel( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + function.dispatch( + options, + args.dtypes(), + ExecuteRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + ctx, + ), + ) +} + +fn try_execute_rows_unfiltered( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + function.dispatch( + options, + args.dtypes(), + ExecuteValidRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + valid, + ctx, + ), + ) +} + +fn prepare_batch( + function: &F, + options: &F::Options, + args: &dyn ExecutionArgs, +) -> VortexResult { + Batch::new(RowFn::id(function), args, |arg_dtypes| { + function.dispatch( + options, + arg_dtypes, + BatchPlanner::::new(arg_dtypes, options), + ) + }) +} + #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use super::execute_rows; use super::row_fn_return_dtype; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::RowVisitor; + use crate::validity::Validity; #[derive(Clone)] struct IndexingRowFn; + #[derive(Clone)] + struct ChangingDispatchRowFn { + dispatches: Arc, + change: DispatchChange, + } + + #[derive(Clone, Copy)] + enum DispatchChange { + Policy, + Element, + } + impl RowFn for IndexingRowFn { type Options = EmptyOptions; @@ -168,6 +272,35 @@ mod tests { } } + impl RowFn for ChangingDispatchRowFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.changing_dispatch_row_fn"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { + visitor.visit::<(i64,), i64>(|(value,)| value) + } else { + match self.change { + DispatchChange::Policy => visitor + .visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())), + DispatchChange::Element => visitor.visit::<(u64,), u64>(|(value,)| value), + } + } + } + } + #[test] fn test_return_dtype_rejects_wrong_arity_before_dispatch() { let error = row_fn_return_dtype(&IndexingRowFn, &EmptyOptions, &[]) @@ -186,6 +319,52 @@ mod tests { assert_arity_error(error); } + #[test] + fn test_execute_rejects_dispatch_that_changes_after_planning() { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Policy, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("dispatch must not change after planning"); + let message = error.to_string(); + + assert!( + message.contains("row dispatch must select the planned nullable execution policy"), + "unexpected error: {error}", + ); + assert!( + message.contains("planned Dense, got DenseWithRetry"), + "unexpected error: {error}", + ); + } + + #[test] + fn test_execute_revalidates_element_types_after_planning() -> VortexResult<()> { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Element, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&function, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_error::vortex_bail!("dispatch must preserve its planned element types"), + }; + + assert!( + error.to_string().contains("expected a u64 column"), + "unexpected error: {error}", + ); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index 98043caec13..fd9685dcb35 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -685,6 +685,14 @@ fn trace_compare_on_dict() -> VortexResult<()> { iter 0 current=vortex.dict(bool, len=5) builder_active=false ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) iter 1 current=vortex.binary(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=5) diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index d779757bc77..bf2b00f563e 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -209,6 +209,14 @@ fn trace_scan_compare_on_compressed_shipdate() -> VortexResult<()> { Done array=vortex.primitive(i32, len=4096) iter 1 current=vortex.primitive(i32, len=4096) builder_active=false return output=vortex.primitive(i32, len=4096) + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=4096) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=4096) iter 2 current=vortex.bool(bool, len=4096) builder_active=false return output=vortex.bool(bool, len=4096) @@ -267,6 +275,14 @@ fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { Done array=vortex.primitive(i16, len=50) iter 1 current=vortex.primitive(i16, len=50) builder_active=false return output=vortex.primitive(i16, len=50) + optimize root=vortex.slice(i16, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i16, len=1) child=vortex.constant(i16, len=50) -> vortex.constant(i16, len=1) + done output=vortex.constant(i16, len=1) + execute_until target=AnyCanonical root=vortex.constant(i16, len=1) + iter 0 current=vortex.constant(i16, len=1) builder_active=false + Done array=vortex.primitive(i16, len=1) + iter 1 current=vortex.primitive(i16, len=1) builder_active=false + return output=vortex.primitive(i16, len=1) Done array=vortex.bool(bool, len=50) iter 6 current=vortex.bool(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=4096) diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index aac8ead42fe..1da0b037f2a 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -479,6 +479,33 @@ impl BitBuffer { } } + /// Fallible variant of [`for_each_set_index`](Self::for_each_set_index). + /// + /// Stops and returns the first error from `f`. + #[inline] + pub fn try_for_each_set_index(&self, mut f: F) -> Result<(), E> + where + F: FnMut(usize) -> Result<(), E>, + { + let mut base = 0usize; + for word in self.chunks().iter_padded() { + if word == u64::MAX { + for k in 0..64 { + f(base + k)?; + } + } else { + let mut w = word; + while w != 0 { + f(base + w.trailing_zeros() as usize)?; + w &= w - 1; + } + } + base += 64; + } + + Ok(()) + } + /// Created a new BitBuffer with offset reset to 0 pub fn sliced(&self) -> Self { if self.offset.is_multiple_of(8) { @@ -970,12 +997,21 @@ mod tests { #[case(65)] #[case(200)] #[case(1000)] - fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + fn test_set_index_visitors_match_set_indices(#[case] len: usize) { let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); buf.for_each_set_index(|i| got.push(i)); assert_eq!(got, expected); + + let mut fallible_got = Vec::new(); + let result = buf.try_for_each_set_index(|i| { + fallible_got.push(i); + Ok::<(), ()>(()) + }); + assert_eq!(result, Ok(())); + assert_eq!(fallible_got, expected); } #[rstest] @@ -998,6 +1034,33 @@ mod tests { assert_eq!(got, (0..130).collect::>()); } + #[test] + fn test_try_for_each_set_index_stops_on_error() { + for (buffer, stop) in [ + (BitBuffer::new_set(130), 65), + (BitBuffer::collect_bool(130, |i| i % 3 == 0), 66), + ] { + let mut visited = Vec::new(); + let result = buffer.try_for_each_set_index(|index| { + visited.push(index); + if index == stop { + return Err(index); + } + + Ok(()) + }); + + assert_eq!(result, Err(stop)); + assert_eq!( + visited, + buffer + .set_indices() + .take_while(|&i| i <= stop) + .collect::>() + ); + } + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 3be2b2d9d66..cd306089325 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -22,7 +22,7 @@ geo-types = { workspace = true } geoarrow = { workspace = true } geoarrow-cast = { workspace = true } prost = { workspace = true } -vortex-array = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-edition = { workspace = true } diff --git a/vortex-spatial/benches/binary_predicates.rs b/vortex-spatial/benches/binary_predicates.rs index b84ab67b17c..281578e4e16 100644 --- a/vortex-spatial/benches/binary_predicates.rs +++ b/vortex-spatial/benches/binary_predicates.rs @@ -12,11 +12,6 @@ //! column-x-column arms are the control: no operand is constant, so a prepared path has nothing to //! hoist and must not regress them. //! -//! `contains` has no all-overlapping arm. One `contains(query polygon, contained square)` row -//! builds a topology graph over the constant's 128 edges, which CodSpeed's CPU simulation charges -//! around 120 µs, so no row count both fits the per-iteration budget and exercises the row loop. -//! [`intersects::polygons_overlapping_x_constant`] covers the never-rejects case instead. -//! //! Run with `cargo bench -p vortex-spatial --bench binary_predicates`. #![expect(clippy::unwrap_used)] @@ -64,6 +59,10 @@ const ROWS: usize = 1 << 7; /// pairwise predicate. It needs a smaller fixture than [`ROWS`] to stay inside the same budget. const OVERLAPPING_POLYGON_ROWS: usize = 1 << 5; +/// Containment builds a topology graph for each polygon pair. Four rows fit the benchmark budget +/// while exercising construction followed by reuse of the prepared constant geometry. +const CONTAINED_POLYGON_ROWS: usize = 4; + /// Deterministic pseudo-random value in `[0, 1)`. fn unit(i: usize) -> f64 { ((i.wrapping_mul(2654435761) >> 8) % 10_000) as f64 / 10_000.0 @@ -225,6 +224,23 @@ mod contains { }); } + /// Constant container against contained polygons: every bbox check passes, the first row + /// prepares the constant geometry, and the remaining rows reuse it for the full predicate. + #[divan::bench] + fn constant_x_polygons_overlapping(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); + let query = query_constant(&mut ctx, CONTAINED_POLYGON_ROWS); + let polygons = squares_mostly_overlapping(CONTAINED_POLYGON_ROWS); + bencher + .counter(ItemsCount::new(CONTAINED_POLYGON_ROWS)) + .bench_local(|| { + execute( + SpatialContains::try_new_array(query.clone(), polygons.clone()), + &mut ctx, + ) + }); + } + /// Constant container against a point column with one null row in eight. #[divan::bench] fn constant_x_nullable_points(bencher: Bencher) { diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index f24f31f02aa..8a670bb0e12 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -178,6 +178,63 @@ pub(crate) fn geometries( } } +/// The geometry a null row decodes to under [`geometries_null_tolerant`]. Arbitrary: the caller +/// guarantees null rows are never read. +pub(crate) fn placeholder_geometry() -> Geometry { + Geometry::Point(geo_types::Point::new(0.0, 0.0)) +} + +/// Whether [`geometries_null_tolerant`] supports this array without filtering null rows first. +pub(crate) fn can_decode_geometries_null_tolerant(array: &ArrayRef) -> VortexResult { + if array.validity()?.definitely_no_nulls() { + return Ok(true); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + + Ok(ext.is::() || ext.is::()) +} + +/// Decode a native geometry column that may contain null rows, writing [`placeholder_geometry`] +/// into their slots. The caller guarantees null rows are never read. +/// +/// `Ok(None)` means this geometry type has no null-tolerant decode yet (`Point` and `Polygon` are +/// covered), and the caller falls back to the filter strategy, which never decodes a null row. A +/// column with definitely no nulls delegates to the ordinary [`geometries`] for any type. +pub(crate) fn geometries_null_tolerant( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>>> { + if array.validity()?.definitely_no_nulls() { + return geometries(array, ctx).map(Some); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + + if ext.is::() { + point_geometries_null_tolerant(&storage, ctx).map(Some) + } else if ext.is::() { + polygon_geometries_null_tolerant(&storage, ctx).map(Some) + } else { + Ok(None) + } +} + /// Decode a constant operand scalar to one geometry, a constant of any /// supported geometry type is decoded exactly like a column. pub(crate) fn single_geometry( diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index e6a00fe8fea..b774f624c4c 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -50,6 +50,7 @@ use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. @@ -150,6 +151,23 @@ pub(crate) fn point_geometries( .collect() } +/// Like [`point_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn point_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + point_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + impl ArrowExportVTable for Point { fn arrow_ext_id(&self) -> Id { *ARROW_POINT diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index a4c88b07b22..bce33efe6e5 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -53,6 +53,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). @@ -153,6 +154,23 @@ pub(crate) fn polygon_geometries( .collect() } +/// Like [`polygon_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn polygon_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + polygon_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + /// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let polygon_type = polygon_type( diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 8850e59f751..9b22a471204 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -3,44 +3,31 @@ //! `ST_Contains`: OGC containment test between two native geometries. +use std::cell::OnceCell; + +use geo::BoundingRect; use geo::Contains; +use geo::PreparedGeometry; +use geo::Relate; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Contains`. -fn validate_contains_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: contains requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: contains operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Contains` between two native geometry operands, each a column or a constant /// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone @@ -59,83 +46,314 @@ impl SpatialContains { } } -impl ScalarFnVTable for SpatialContains { +impl RowFn for SpatialContains { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.contains"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) + /// Containment is not symmetric, so `a` is always the container and `b` the contained. + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), UninitElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstOperands { + a: a.map(PreparedOperand::new), + b: b.map(PreparedOperand::new), + } + }, + |operands, (a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, contains_row_prepared(operands, a, b)) } + }, + ) } +} + +/// Per-batch state for the contains row kernel: the prepared form of whichever operand is +/// constant for the batch. `None` marks an operand that varies by row. +struct ConstOperands { + /// Operand `a` (the container) when it is batch-constant. + a: Option, + + /// Operand `b` (the contained) when it is batch-constant. + b: Option, +} + +/// One batch-constant operand: its bounding rectangle and the [`PreparedGeometry`] built on the +/// first row whose pairing routes through relate. +/// +/// The build is lazy because preparation (self-noding the topology graph plus an R*-tree over the +/// edges) costs `O(edges log edges)` and pays off only on relate-routed pairings; a batch of +/// point rows against a constant polygon never touches it, and preparing a large constant eagerly +/// would charge such a batch for nothing. +struct PreparedOperand { + /// The constant's bounding rectangle, folded once for conservative row rejection. + bbox: Option>, + + /// The constant's prepared form, initialized only when a relate route needs it. + prepared: OnceCell, f64>>, +} - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("contains has exactly two children"), +impl PreparedOperand { + fn new(geometry: &Geometry) -> Self { + Self { + bbox: finite_bounding_rect(geometry), + prepared: OnceCell::new(), } } - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_contains_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) + /// Return the prepared geometry, cloning the decoded constant only on first use. + /// + /// `geometry` **must** be the constant represented by this state. The row kernel maintains + /// that relationship by passing the operand from the same decoded constant column that + /// produced this [`PreparedOperand`]. + fn get(&self, geometry: &Geometry) -> &PreparedGeometry<'static, Geometry, f64> { + self.prepared + .get_or_init(|| PreparedGeometry::from(geometry.clone())) } +} - fn execute( - &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Containment is not symmetric: `a` is always the container and `b` the contained. A - // container's rect must cover the contained's rect (`Rect::contains` is the closed - // test), so a contained rect poking outside proves the row false. - execute_binary_geo_types( - &a, - &b, - |a, b| a.contains(b), - Some(|ra, rb| (!ra.contains(rb)).then_some(false)), - ctx, - ) - } +/// Returns a bounding rectangle only when ordered comparisons can conservatively reject a row. +/// +/// Geo permits non-finite coordinates. A rectangle containing NaN cannot prove non-containment, +/// because its ordered comparisons can return false even when the exact algorithm accepts the +/// geometry. +fn finite_bounding_rect(geometry: &Geometry) -> Option> { + let bbox = geometry.bounding_rect()?; + let min = bbox.min(); + let max = bbox.max(); + + [min.x, min.y, max.x, max.y] + .into_iter() + .all(f64::is_finite) + .then_some(bbox) +} - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) +/// How geo's `a.contains(b)` computes its verdict for a pairing. +enum ContainsRoute { + /// `a.relate(b).is_contains()`. + ForwardRelate, + + /// `b.relate(a).is_within()`, how geo phrases relate for `MultiPolygon` containers. + ReversedRelate, + + /// A direct algorithm (coordinate position, point arithmetic); nothing to prepare. + Direct, +} + +/// The route geo 0.31's `Contains` dispatch takes for `a.contains(b)`. +/// +/// The prepared substitution in [`contains_row_prepared`] **must** run relate exactly where geo +/// runs relate, with the same argument order, because geo's direct algorithms are not everywhere +/// bit-identical to a relate matrix query (they resolve degenerate and boundary cases with +/// different arithmetic). The relate rows below transcribe geo's `impl_contains_from_relate!` +/// lists per container type; everything else, notably every `Point`/`MultiPoint` contained side +/// and every `Point` container, is direct. +/// +/// **This table is coupled to the geo version.** It transcribes a dispatch that geo is free to +/// reshuffle in any release, and a wrong row is a silently wrong verdict rather than a build error. +/// The workspace therefore pins `geo = "=0.31.0"`: taking any new geo, patch releases included, is +/// a deliberate edit of that line, and the edit must re-verify this table against +/// `impl_contains_from_relate!`. +/// +/// `constant_operands_agree_with_columns` is the mechanical check, and it is **not** complete: it +/// compares the prepared route against plain `a.contains(b)` only for the container types it has +/// cases for. `routes_agree_with_geo_for_every_container` covers the rest, one representative +/// pairing per container variant, and is the one to extend when geo grows a geometry type. Both +/// stay green wherever relate and the direct algorithm agree, so neither replaces the pin. +fn contains_route(a: &Geometry, b: &Geometry) -> ContainsRoute { + use Geometry as G; + + match (a, b) { + // Line contains [Polygon, MultiLineString, MultiPolygon, GeometryCollection, Rect, + // Triangle]. + ( + G::Line(_), + G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // LineString contains [Polygon, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::LineString(_), + G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiLineString contains everything except Point. + | ( + G::MultiLineString(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiPoint contains [Line, LineString, Polygon, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::MultiPoint(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Polygon contains everything except Point and MultiPoint. + | ( + G::Polygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Rect contains [Line, LineString, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Triangle]; Rect contains Rect and Polygon are direct. + | ( + G::Rect(_), + G::Line(_) + | G::LineString(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Triangle(_), + ) + // Triangle and GeometryCollection contain everything except Point. + | ( + G::Triangle(_) | G::GeometryCollection(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ForwardRelate, + + // MultiPolygon contains everything except Point and MultiPoint, phrased reversed. + ( + G::MultiPolygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ReversedRelate, + + _ => ContainsRoute::Direct, } +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +/// Computes one row of contains, substituting a prepared graph for a constant operand on the +/// pairings geo itself answers through relate. +/// +/// [`PreparedGeometry`] carries the operand's self-noded topology graph and edge R*-tree, so a +/// relate against it skips rebuilding both and reads its bounding rect from cache; geo asserts +/// the cached graph equal to a freshly built one (its `swap_arg_index` test), which is what makes +/// the substitution result-preserving. Before dispatch, a disjoint constant-side bounding rect +/// conservatively rejects the row, matching the columnar implementation's #9076 optimization. +/// All other rows delegate to the same direct or relate route as `a.contains(b)`. +fn contains_row_prepared(operands: &ConstOperands, a: &Geometry, b: &Geometry) -> bool { + let rejected = match (&operands.a, &operands.b) { + (None, None) => false, + (Some(const_a), Some(const_b)) => const_a + .bbox + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (Some(const_a), None) => const_a + .bbox + .zip(finite_bounding_rect(b)) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (None, Some(const_b)) => finite_bounding_rect(a) + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + }; + + if rejected { + return false; } - fn is_fallible(&self, _: &Self::Options) -> bool { - false + match contains_route(a, b) { + ContainsRoute::Direct => a.contains(b), + ContainsRoute::ForwardRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_a.get(a).relate(const_b.get(b)).is_contains(), + (Some(const_a), None) => const_a.get(a).relate(b).is_contains(), + (None, Some(const_b)) => a.relate(const_b.get(b)).is_contains(), + (None, None) => a.contains(b), + }, + ContainsRoute::ReversedRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_b.get(b).relate(const_a.get(a)).is_within(), + (Some(const_a), None) => b.relate(const_a.get(a)).is_within(), + (None, Some(const_b)) => const_b.get(b).relate(a).is_within(), + (None, None) => a.contains(b), + }, } } #[cfg(test)] mod tests { + use geo::Contains; + use geo_types::Coord; use geo_types::Geometry; + use geo_types::GeometryCollection; + use geo_types::Line; use geo_types::LineString; + use geo_types::MultiLineString; + use geo_types::MultiPoint; + use geo_types::MultiPolygon; use geo_types::Point; use geo_types::Polygon; + use geo_types::Rect; + use geo_types::Triangle; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -144,6 +362,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::MaskedArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -158,10 +377,15 @@ mod tests { use vortex_error::vortex_err; use wkb::writer::WriteOptions; + use super::ConstOperands; + use super::PreparedOperand; use super::SpatialContains; + use super::contains_row_prepared; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::linestring_column; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::polygon_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -218,6 +442,21 @@ mod tests { assert_contains(container, other, [expected; 3]) } + /// A non-finite bounding rectangle cannot reject a containment that the exact geometry + /// algorithm accepts. + #[test] + fn nan_bounding_rect_does_not_reject_containment() { + let container = multipoint(vec![(f64::NAN, f64::NAN), (1.0, 1.0)]); + let contained = point(1.0, 1.0); + let operands = ConstOperands { + a: Some(PreparedOperand::new(&container)), + b: Some(PreparedOperand::new(&contained)), + }; + + assert!(container.contains(&contained)); + assert!(contains_row_prepared(&operands, &container, &contained)); + } + /// Partially overlapping polygons contain each other in neither direction. #[test] fn overlapping_polygons_contain_neither_way() -> VortexResult<()> { @@ -246,6 +485,20 @@ mod tests { assert_contains(container, points, [true, false, false]) } + /// Constant container vs a linestring column: a row whose bounding rect pokes outside the + /// container's is not contained, while one wholly inside is. Carried over from the columnar + /// bounding-rect rejection in #9076, since it constrains the verdict rather than the mechanism. + #[test] + fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let lines = linestring_column(vec![ + vec![(1.0, 1.0), (3.0, 3.0)], + vec![(1.0, 1.0), (9.0, 1.0)], + vec![(5.0, 5.0), (9.0, 9.0)], + ])?; + assert_contains(container, lines, [true, false, false]) + } + /// Polygon column vs constant point: only the polygon around the point contains it. #[test] fn polygon_column_vs_constant_point() -> VortexResult<()> { @@ -266,20 +519,6 @@ mod tests { assert_contains(away, point, [false; 2]) } - /// Constant container vs a linestring column: a row whose bounding rect pokes outside the - /// container's rect is proven false by the rect pre-check alone; a fully inside row still - /// needs (and passes) the exact test. - #[test] - fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { - let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; - let lines = linestring_column(vec![ - vec![(1.0, 1.0), (3.0, 3.0)], - vec![(1.0, 1.0), (9.0, 1.0)], - vec![(5.0, 5.0), (9.0, 9.0)], - ])?; - assert_contains(container, lines, [true, false, false]) - } - /// Column vs column pairs rows: each polygon row is tested against the point row at the /// same position. #[test] @@ -410,6 +649,83 @@ mod tests { Ok(()) } + /// A nullable polygon column: unit squares at `centers`, the rows where `nulls` is true + /// masked out, spelled as `Masked` over non-nullable storage. + fn nullable_squares(centers: &[(f64, f64)], nulls: &[bool]) -> VortexResult { + let squares = centers + .iter() + .map(|&(x, y)| { + vec![vec![ + (x - 1.0, y - 1.0), + (x + 1.0, y - 1.0), + (x + 1.0, y + 1.0), + (x - 1.0, y + 1.0), + (x - 1.0, y - 1.0), + ]] + }) + .collect(); + let polygons = polygon_column(squares)?; + + Ok( + MaskedArray::try_new(polygons, Validity::from_iter(nulls.iter().map(|n| !n)))? + .into_array(), + ) + } + + /// Nullable geometry operands conjoin their validity before computing containment. + #[test] + fn contains_nullable_geometries_conjoins_validity() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let centers = [(0.0, 0.0), (5.0, 5.0), (0.5, -0.2), (9.0, 9.0), (0.0, 1.0)]; + let nulls = [false, true, false, false, true]; + let polygons = nullable_squares(¢ers, &nulls)?; + let points = nullable_point_column(vec![ + Some((0.0, 0.0)), + Some((5.0, 5.0)), + None, + Some((0.0, 0.0)), + Some((0.0, 1.0)), + ])?; + + let actual = SpatialContains::try_new_array(polygons, points)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + let expected = BoolArray::from_iter([Some(true), None, None, Some(false), None]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// Geometry types without a null-tolerant decode fall back to filtering valid rows. + #[test] + fn contains_unsupported_geometry_falls_back_to_filter() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let validity = Validity::from_iter([true, false, true, true]); + let lines = linestring_column(vec![ + vec![(0.0, 0.0), (4.0, 4.0)], + vec![(0.0, 0.0), (1.0, 1.0)], + vec![(2.0, 2.0), (3.0, 3.0)], + vec![(0.0, 4.0), (4.0, 0.0)], + ])?; + let nullable_lines = MaskedArray::try_new(lines.clone(), validity.clone())?.into_array(); + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 4)?; + + let expected = SpatialContains::try_new_array(lines, point.clone())?.into_array(); + let expected = MaskedArray::try_new(expected, validity)?.into_array(); + let actual = SpatialContains::try_new_array(nullable_lines, point)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + /// A non-geometry operand dtype is rejected up front, before execution. #[test] fn non_geometry_operand_is_rejected() -> VortexResult<()> { @@ -419,4 +735,166 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// A two-point line segment geometry, the `Line` container variant. + fn line_geometry(start: (f64, f64), end: (f64, f64)) -> Geometry { + Geometry::Line(Line::new( + Coord { + x: start.0, + y: start.1, + }, + Coord { x: end.0, y: end.1 }, + )) + } + + /// A multilinestring geometry over one linestring per entry of `parts`. + fn multilinestring(parts: Vec>) -> Geometry { + Geometry::MultiLineString(MultiLineString::new( + parts.into_iter().map(LineString::from).collect(), + )) + } + + /// A geometry collection wrapping `parts`. + fn collection(parts: Vec) -> Geometry { + Geometry::GeometryCollection(GeometryCollection::from(parts)) + } + + /// An axis-aligned rectangle geometry, the `Rect` container variant. + fn rect_geometry(x0: f64, y0: f64, x1: f64, y1: f64) -> Geometry { + Geometry::Rect(Rect::new(Coord { x: x0, y: y0 }, Coord { x: x1, y: y1 })) + } + + /// A triangle geometry large enough to contain the small test polygons. + fn triangle_geometry() -> Geometry { + Geometry::Triangle(Triangle::new( + Coord { x: 0.0, y: 0.0 }, + Coord { x: 8.0, y: 0.0 }, + Coord { x: 0.0, y: 8.0 }, + )) + } + + /// A two-part multipolygon: `4x4` squares at the origin and at `(10, 10)`. + fn two_part_multipolygon() -> Geometry { + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 4.0, 4.0), + rect_polygon(10.0, 10.0, 14.0, 14.0), + ])) + } + + /// Every container variant `contains_route` distinguishes, checked against plain + /// `a.contains(b)` in all four constant arrangements. + /// + /// Every case is a containment geo answers `true`, which the test asserts: a pairing that is + /// false regardless of route (a lower-dimensional container, say) also agrees regardless of + /// route, and pins nothing. A true case fails when the prepared substitution diverges from + /// geo: a table row whose relate phrasing disagrees with geo's dispatch on this input, or a + /// bounding-rect prescreen that wrongly rejects a contained row. It is **not** a version + /// tripwire: a geo release that reshuffles its dispatch stays green wherever relate and the + /// direct algorithm agree, which is why the workspace pins `geo` exactly. + /// + /// This is the table's own regression, and the one to extend when geo grows a geometry type: + /// `constant_operands_agree_with_columns` below goes through real arrays and so is the better + /// end-to-end check, but it only covers the container types it has cases for, and WKB decoding + /// limits which types those can be. The MultiPoint and Line containers route relate only for + /// contained types a MultiPoint or Line can rarely contain, so their true cases lean on + /// `GeometryCollection` membership and collinear `MultiLineString` parts respectively. + #[rstest] + #[case::point(point(1.0, 1.0), point(1.0, 1.0))] + #[case::line(line_geometry((0.0, 0.0), (4.0, 4.0)), point(2.0, 2.0))] + #[case::line_x_multilinestring(line_geometry((0.0, 0.0), (4.0, 4.0)), multilinestring(vec![vec![(1.0, 1.0), (2.0, 2.0)]]))] + #[case::linestring(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::multipoint(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), collection(vec![point(2.0, 2.0)]))] + #[case::multilinestring(multilinestring(vec![vec![(0.0, 0.0), (4.0, 4.0)]]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::geometrycollection(collection(vec![rect_polygon(0.0, 0.0, 8.0, 8.0).into()]), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::rect(rect_geometry(0.0, 0.0, 8.0, 8.0), line(vec![(2.0, 2.0), (4.0, 4.0)]))] + #[case::triangle(triangle_geometry(), rect_polygon(1.0, 1.0, 2.0, 2.0).into())] + fn routes_agree_with_geo_for_every_container(#[case] a: Geometry, #[case] b: Geometry) { + let expected = a.contains(&b); + assert!( + expected, + "route cases must be containments geo answers true, or every route agrees vacuously", + ); + + let arrangements = [ + (None, None), + (Some(PreparedOperand::new(&a)), None), + (None, Some(PreparedOperand::new(&b))), + ( + Some(PreparedOperand::new(&a)), + Some(PreparedOperand::new(&b)), + ), + ]; + + for (index, (const_a, const_b)) in arrangements.into_iter().enumerate() { + let operands = ConstOperands { + a: const_a, + b: const_b, + }; + assert_eq!( + contains_row_prepared(&operands, &a, &b), + expected, + "arrangement {index} disagrees with geo's own contains", + ); + } + } + + /// Constant arrangements agree with expanded columns across the routes the prepared kernel + /// distinguishes: forward relate (polygon, linestring and multipoint containers), reversed + /// relate (multipolygon containers), and the direct pairings (a point on either side, + /// multipoint over multipoint, polygon over multipoint), including boundary contact, + /// crossing, disjoint and empty cases. + #[rstest] + #[case::polygon_nested_polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_touching_from_inside(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(0.0, 2.0, 2.0, 4.0).into())] + #[case::polygon_overlapping_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygon_disjoint_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygon_x_point_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(0.0, 2.0))] + #[case::polygon_x_point_outside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(20.0, 20.0))] + #[case::polygon_x_nan_point(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(f64::NAN, 2.0))] + #[case::polygon_x_linestring_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_linestring_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::polygon_x_linestring_crossing(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(-2.0, 2.0), (2.0, 2.0)]))] + #[case::polygon_x_empty_linestring(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![]))] + #[case::polygon_x_multipoint_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_multipoint_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::linestring_x_multipoint_on_line(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipoint_x_multipoint_subset(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), multipoint(vec![(2.0, 2.0)]))] + #[case::multipoint_x_linestring_between_points(multipoint(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon_x_polygon_in_one_part(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::multipolygon_x_polygon_straddling(two_part_multipolygon(), rect_polygon(3.0, 3.0, 11.0, 11.0).into())] + #[case::multipolygon_x_polygon_disjoint(two_part_multipolygon(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::multipolygon_x_point_inside(two_part_multipolygon(), point(11.0, 11.0))] + #[case::point_x_point_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::point_x_polygon(point(2.0, 2.0), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialContains::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index dfd3d09ed23..e41999338a6 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -6,43 +6,20 @@ use geo::Distance; use geo::Euclidean; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Distance`. -fn validate_distance_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: distance requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: distance operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -60,66 +37,41 @@ impl SpatialDistance { } } -impl ScalarFnVTable for SpatialDistance { +impl RowFn for SpatialDistance { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.distance"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { - Ok(EmptyOptions) - } - - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("distance has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_distance_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Primitive(PType::F64, nullability)) - } - - fn execute( + fn deserialize( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Distance is a value, not a verdict: no bounding-rect test can decide it. - execute_binary_geo_types(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx) + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn validity( + fn dispatch>( &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } - - fn is_strict(&self, _: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _: &Self::Options) -> bool { - false + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(GeometryRow, GeometryRow), UninitElementSink, _>( + |(a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, Euclidean.distance(a, b)) } + }, + ) } } @@ -196,8 +148,9 @@ mod tests { Ok(()) } - /// Distance passes no bounding-rect rejection: a point far outside a constant polygon's - /// bounding rect still gets its true distance, alongside an inside point at distance zero. + /// Distance is a value rather than a verdict, so no bounding-rect rejection may fire for it: a + /// point far outside a constant polygon's rect still gets its true distance. Carried over from + /// #9076, which added the rejection to the predicates but deliberately not to this function. #[test] fn distance_to_constant_polygon_is_exact() -> VortexResult<()> { let session = vortex_array::array_session(); diff --git a/vortex-spatial/src/scalar_fn/execute.rs b/vortex-spatial/src/scalar_fn/execute.rs index 3a7494bcb39..2acdce76e36 100644 --- a/vortex-spatial/src/scalar_fn/execute.rs +++ b/vortex-spatial/src/scalar_fn/execute.rs @@ -7,17 +7,14 @@ //! propagation without prescribing how a kernel represents geometries or builds its output. //! Native columnar kernels such as `ST_MakeLine` use these dispatchers directly. //! -//! [`execute_unary_geo_types`] and [`execute_binary_geo_types`] are convenience adapters for -//! row-oriented algorithms from the `geo` ecosystem. They decode valid inputs into -//! `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`], such as an `f64` or -//! boolean array. +//! [`execute_unary_geo_types`] adapts row-oriented algorithms from the `geo` ecosystem. It decodes +//! valid inputs into `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`]. mod binary; mod geo_types; mod unary; pub(crate) use binary::dispatch_binary; -pub(crate) use binary::execute_binary_geo_types; pub(crate) use unary::dispatch_unary; pub(crate) use unary::execute_unary_geo_types; use vortex_array::ArrayRef; diff --git a/vortex-spatial/src/scalar_fn/execute/binary.rs b/vortex-spatial/src/scalar_fn/execute/binary.rs index f2c03bd1beb..5cf639461b0 100644 --- a/vortex-spatial/src/scalar_fn/execute/binary.rs +++ b/vortex-spatial/src/scalar_fn/execute/binary.rs @@ -1,28 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Binary pairwise dispatch, plus an adapter for row-oriented `geo_types` kernels. +//! Binary constant-and-column operand dispatch. -use geo::BoundingRect; -use geo_types::Geometry; -use geo_types::Rect; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_mask::Mask; use super::Execution; use super::Operand; -use super::geo_types::GeoTypesOutput; -use super::geo_types::eval_column; -use super::geo_types::eval_column_pair; -use crate::extension::single_geometry; /// Dispatch a binary strict geometry kernel over constants and columns. /// @@ -80,6 +72,7 @@ where if len != 0 && valid.all_false() { return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); } + kernel( Execution { operands: [left, right], @@ -90,245 +83,3 @@ where ctx, ) } - -/// A bounding-rectangle pre-check for [`execute_binary_geo_types`]'s one-constant paths. -/// -/// Called per row with rectangles in operand order, it returns `Some(result)` when they prove the -/// result and `None` when the exact kernel must run. -pub(crate) type BboxPrecheck = fn(&Rect, &Rect) -> Option; - -/// Run a binary row-oriented kernel whose inputs are decoded to `geo_types::Geometry`. -/// -/// The `geo_types` name describes the values passed to `compute`, not the output. `T` is converted -/// into a Vortex array before this function returns. Nulls propagate from either operand. With -/// exactly one constant operand, `bbox_precheck` may prove a result from the fixed constant -/// bounding rectangle and the current row's rectangle before the exact kernel runs. -pub(crate) fn execute_binary_geo_types( - left: &ArrayRef, - right: &ArrayRef, - compute: F, - bbox_precheck: Option>, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T + Copy, -{ - let nullability = Nullability::from(left.dtype().is_nullable() || right.dtype().is_nullable()); - dispatch_binary( - left, - right, - T::dtype(nullability), - |execution, ctx| match execution.operands { - [Operand::Constant(left), Operand::Constant(right)] => { - let left = single_geometry(&left, ctx)?; - let right = single_geometry(&right, ctx)?; - Ok(ConstantArray::new( - compute(&left, &right).into_scalar(execution.nullability), - execution.len, - ) - .into_array()) - } - [Operand::Constant(left), Operand::Column(right)] => { - let left = single_geometry(&left, ctx)?; - let prescreen = bbox_precheck.zip(left.bounding_rect()); - eval_column( - &right, - &execution.valid, - |right| { - prescreen - .and_then(|(precheck, fixed)| precheck(&fixed, &right.bounding_rect()?)) - .unwrap_or_else(|| compute(&left, right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Constant(right)] => { - let right = single_geometry(&right, ctx)?; - let prescreen = bbox_precheck.zip(right.bounding_rect()); - eval_column( - &left, - &execution.valid, - |left| { - prescreen - .and_then(|(precheck, fixed)| precheck(&left.bounding_rect()?, &fixed)) - .unwrap_or_else(|| compute(left, &right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Column(right)] => eval_column_pair( - &left, - &right, - &execution.valid, - compute, - execution.nullability, - ctx, - ), - }, - ctx, - ) -} - -#[cfg(test)] -mod tests { - use std::cell::Cell; - - use geo::Contains; - use geo::Intersects; - use geo_types::Geometry; - use vortex_array::ArrayRef; - use vortex_array::ExecutionCtx; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ConstantArray; - use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; - use vortex_buffer::BitBuffer; - use vortex_error::VortexResult; - - use super::BboxPrecheck; - use super::execute_binary_geo_types; - use crate::test_harness::linestring_column; - use crate::test_harness::nullable_point_column; - use crate::test_harness::point_column; - use crate::test_harness::polygon_column; - - const DISJOINT_PRECHECK: BboxPrecheck = - |left, right| (!left.intersects(right)).then_some(false); - - fn triangle_constant(len: usize, ctx: &mut ExecutionCtx) -> VortexResult { - let ring = vec![(0.0, 0.0), (10.0, 0.0), (0.0, 10.0), (0.0, 0.0)]; - let scalar = polygon_column(vec![vec![ring]])?.execute_scalar(0, ctx)?; - Ok(ConstantArray::new(scalar, len).into_array()) - } - - fn counting_intersects( - counter: &Cell, - ) -> impl Fn(&Geometry, &Geometry) -> bool + Copy { - move |left, right| { - counter.set(counter.get() + 1); - left.intersects(right) - } - } - - #[test] - fn bbox_precheck_skips_exact_test() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = point_column(vec![50.0, 8.0, 2.0], vec![50.0, 8.0, 2.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_leaves_nulls_alone() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = nullable_point_column(vec![Some((50.0, 50.0)), None, Some((2.0, 2.0))])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - let expected = BoolArray::new( - BitBuffer::from_iter([false, false, true]), - Validity::from_iter([true, false, true]), - ) - .into_array(); - - assert_arrays_eq!(result, expected, &mut ctx); - assert_eq!(exact_runs.get(), 1); - Ok(()) - } - - #[test] - fn bbox_precheck_sees_rects_in_operand_order() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let triangle = triangle_constant(2, &mut ctx)?; - let exact_runs = Cell::new(0); - let counted = |left: &Geometry, right: &Geometry| { - exact_runs.set(exact_runs.get() + 1); - left.contains(right) - }; - - let result = execute_binary_geo_types( - &probes, - &triangle, - counted, - Some(|left, right| (!left.contains(right)).then_some(false)), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 0); - Ok(()) - } - - #[test] - fn empty_constant_falls_through_to_exact() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let scalar = linestring_column(vec![vec![]])?.execute_scalar(0, &mut ctx)?; - let empty = ConstantArray::new(scalar, 2).into_array(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &empty, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_matches_exact_results() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(6, &mut ctx)?; - let probes = nullable_point_column(vec![ - Some((50.0, 50.0)), - Some((8.0, 8.0)), - Some((2.0, 2.0)), - None, - Some((0.0, 0.0)), - Some((10.0, 0.0)), - ])?; - let exact = |left: &Geometry, right: &Geometry| left.intersects(right); - - let with_precheck = - execute_binary_geo_types(&triangle, &probes, exact, Some(DISJOINT_PRECHECK), &mut ctx)?; - let exact_only = execute_binary_geo_types(&triangle, &probes, exact, None, &mut ctx)?; - - assert_arrays_eq!(with_precheck, exact_only, &mut ctx); - Ok(()) - } -} diff --git a/vortex-spatial/src/scalar_fn/execute/geo_types.rs b/vortex-spatial/src/scalar_fn/execute/geo_types.rs index 038aca46502..7007f02cfc6 100644 --- a/vortex-spatial/src/scalar_fn/execute/geo_types.rs +++ b/vortex-spatial/src/scalar_fn/execute/geo_types.rs @@ -118,27 +118,3 @@ where let values = decoded.iter().map(compute).collect(); Ok(T::build_array(len, valid, values, nullability)) } - -/// Evaluate a decoded kernel over rows where both geometry columns are valid. -pub(super) fn eval_column_pair( - left: &ArrayRef, - right: &ArrayRef, - valid: &Mask, - compute: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T, -{ - let len = left.len(); - let left = geometries(&left.filter(valid.clone())?, ctx)?; - let right = geometries(&right.filter(valid.clone())?, ctx)?; - let values = left - .iter() - .zip(&right) - .map(|(left, right)| compute(left, right)) - .collect(); - Ok(T::build_array(len, valid, values, nullability)) -} diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index 77d33886ff3..9a3a198e838 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -3,44 +3,27 @@ //! `ST_Intersects`: OGC intersection test between two native geometries. +use geo::BoundingRect; use geo::Intersects; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Intersects`. -fn validate_intersects_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: intersects requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: intersects operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry /// operands, each a column or a constant literal. @@ -58,74 +41,100 @@ impl SpatialIntersects { } } -impl ScalarFnVTable for SpatialIntersects { +impl RowFn for SpatialIntersects { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.intersects"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("intersects has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_intersects_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) - } - - fn execute( + fn dispatch>( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Disjoint bounding rects prove the geometries disjoint; rect contact (closed test) - // falls through to the exact test. - execute_binary_geo_types( - &a, - &b, - |x, y| x.intersects(y), - Some(|ra, rb| (!ra.intersects(rb)).then_some(false)), - ctx, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), UninitElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstBboxes::new(a, b) + }, + |bboxes, (a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, intersects_row_prepared(bboxes, a, b)) } + }, ) } +} - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } +/// Per-batch state for the intersects row kernel: the bounding rect of each operand that is +/// constant for the batch. +/// +/// geo opens many intersects pairings with `has_disjoint_bboxes`, an early-out that folds +/// [`bounding_rect`] over both operands. For a batch-constant operand that fold recomputes the +/// same rect every row, so it is hoisted here and [`intersects_row_prepared`] replays the +/// comparison with the hoisted value. `None` marks an operand that varies by row or has no +/// bounding rect (an empty geometry); both mean no early-out, exactly as `has_disjoint_bboxes` +/// treats a missing rect. +/// +/// [`bounding_rect`]: BoundingRect::bounding_rect +struct ConstBboxes { + /// The bounding rect of operand `a` when it is batch-constant. + a: Option>, + + /// The bounding rect of operand `b` when it is batch-constant. + b: Option>, +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +impl ConstBboxes { + fn new(a: Option<&Geometry>, b: Option<&Geometry>) -> Self { + Self { + a: a.and_then(BoundingRect::bounding_rect), + b: b.and_then(BoundingRect::bounding_rect), + } } +} - fn is_fallible(&self, _: &Self::Options) -> bool { - false - } +/// Computes one row of intersects, spending any bounding rect hoisted into `bboxes`. +/// +/// Disjoint bounding rectangles conservatively prove that the geometries do not intersect. The +/// fall-through delegates to the unchanged `a.intersects(b)`, which refolds both rects internally, +/// so a batch where every row overlaps pays one extra `bounding_rect` fold over the row operand; +/// the win concentrates where most rows are disjoint, the usual spatial-filter shape. +fn intersects_row_prepared(bboxes: &ConstBboxes, a: &Geometry, b: &Geometry) -> bool { + let disjoint = match (bboxes.a, bboxes.b) { + (None, None) => false, + (Some(bbox_a), Some(bbox_b)) => !bbox_a.intersects(&bbox_b), + (Some(bbox_a), None) => b + .bounding_rect() + .is_some_and(|bbox_b| !bbox_a.intersects(&bbox_b)), + (None, Some(bbox_b)) => a + .bounding_rect() + .is_some_and(|bbox_a| !bbox_a.intersects(&bbox_b)), + }; + + if disjoint { + return false; + } + + a.intersects(b) } #[cfg(test)] @@ -133,7 +142,9 @@ mod tests { use geo_types::Coord; use geo_types::Geometry; use geo_types::LineString; + use geo_types::MultiPoint; use geo_types::MultiPolygon; + use geo_types::Point; use geo_types::Polygon; use rstest::rstest; use vortex_array::ArrayRef; @@ -158,8 +169,10 @@ mod tests { use wkb::writer::WriteOptions; use super::SpatialIntersects; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::rect_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -441,4 +454,85 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// Constant arrangements agree with expanded columns across the pairing classes the prepared + /// kernel treats differently: bbox-prechecked pairs (polygon x polygon, linestring x + /// anything, multipolygon blankets), direct pairs (points), the excluded `MultiPoint` route, + /// and an empty geometry whose bounding rect does not exist. + #[rstest] + #[case::polygons_overlapping(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygons_touching_edge(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 0.0, 8.0, 4.0).into())] + #[case::polygons_touching_corner(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 4.0, 8.0, 8.0).into())] + #[case::polygons_disjoint(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygons_nested(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_x_point_inside(donut(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(donut(), point(0.0, 5.0))] + #[case::polygon_x_point_in_hole(donut(), point(5.0, 5.0))] + #[case::point_outside_x_polygon(point(20.0, 20.0), donut())] + #[case::nan_point_x_polygon(point(f64::NAN, 2.0), donut())] + #[case::polygon_x_nan_point(donut(), point(f64::NAN, 2.0))] + #[case::points_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::points_distinct(point(1.0, 1.0), point(2.0, 1.0))] + #[case::linestring_crossing_polygon(line(vec![(-2.0, -2.0), (2.0, 2.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_disjoint_polygon(line(vec![(-2.0, -2.0), (-6.0, -6.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_in_polygon_hole(line(vec![(4.5, 4.5), (5.5, 5.5)]), donut())] + #[case::linestrings_crossing(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(0.0, 4.0), (4.0, 0.0)]))] + #[case::linestrings_disjoint(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(10.0, 10.0), (14.0, 14.0)]))] + #[case::empty_linestring_x_polygon(line(vec![]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_straddling_polygon(multipoint(vec![(2.0, 2.0), (20.0, 20.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_outside_polygon(multipoint(vec![(20.0, 20.0), (30.0, 30.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipolygon_disjoint_polygon( + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 2.0, 2.0), + rect_polygon(10.0, 10.0, 12.0, 12.0), + ])), + rect_polygon(20.0, 20.0, 24.0, 24.0).into() + )] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } + + /// `Rect` has no WKB form, so its constant comes from a one-row rect column; its conservative + /// bbox early-out and exact fall-through must agree with the expanded form like the rest. + #[test] + fn rect_operand_agrees_with_columns() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let rect_scalar = rect_column(vec![(0.0, 0.0, 4.0, 4.0)])?.execute_scalar(0, &mut ctx)?; + let rect_constant = ConstantArray::new(rect_scalar, 3).into_array(); + let polygon_constant = + geometry_constant(&Geometry::Polygon(rect_polygon(2.0, 2.0, 6.0, 6.0)), 3)?; + + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + rect_constant, + polygon_constant, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index 99fe5d28528..6291075246a 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -13,3 +13,4 @@ mod execute; pub mod intersects; pub mod length; pub mod make_line; +pub(crate) mod row; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs new file mode 100644 index 00000000000..750497e5f32 --- /dev/null +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the geo scalar functions add to the row-function machinery: an element type that decodes a +//! native geometry column into `geo_types` geometries. + +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::unstable::row::InputElement; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::extension::can_decode_geometries_null_tolerant; +use crate::extension::geometries; +use crate::extension::geometries_null_tolerant; +use crate::extension::is_native_geometry; + +/// Marker for native geometry input elements: accepts any native geometry column and presents each +/// row as a decoded `geo_types` geometry. +/// +/// The two operands of a binary geo function need not share a geometry type, since distance, +/// containment and intersection across types are all meaningful, so this validates only that the +/// column is _some_ native geometry. +pub(crate) struct GeometryRow; + +// SAFETY: [`view`](InputElement::view) returns the decoded geometry slice and +// [`view_len`](InputElement::view_len) reports that slice's exact length. +unsafe impl InputElement for GeometryRow { + type Column = Vec>; + type View<'a> = &'a [Geometry]; + type Elem<'a> = &'a Geometry; + + // A geometry row is decoded from its coordinate storage, which behind a null row holds + // arbitrary coordinates that need not describe a well-formed geometry. + const DENSE_SAFE: bool = false; + // Decoding builds a geometry from stored coordinates, and a malformed one in a _valid_ row is a + // domain error rather than an infrastructural failure. + const DECODE_FALLIBLE: bool = true; + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + is_native_geometry(dtype), + "spatial: operand {dtype} is not a native geometry type" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + geometries(&array, ctx) + } + + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + can_decode_geometries_null_tolerant(array) + } + + fn get(column: &Self::Column, index: usize) -> &Geometry { + &column[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.len() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + &view[index] + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + // SAFETY: The caller established that `index` is below the slice length returned by + // `view_len` for this exact view. + unsafe { view.get_unchecked(index) } + } + + /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. + /// `Point` and `Polygon` columns are covered; other geometry types return `Ok(None)` and the + /// batch falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + geometries_null_tolerant(&array, ctx) + } +} + +/// Test-only support for the prepared geo row kernels: a probe recording which operands a +/// `prepare` step saw as batch-constant, and the shared prepared-vs-expanded agreement check +/// built on it. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::ScalarFnArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + + thread_local! { + /// Which operands the last `prepare` saw as constant, as a bitmask (bit 0 for `a`, bit 1 + /// for `b`). Thread-local rather than a process global so concurrent tests in one process + /// (plain `cargo test`) cannot race it; execution runs on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Record which operands `prepare` saw as constant. + pub(crate) fn record(a_constant: bool, b_constant: bool) { + SEEN_CONSTANTS.set(u8::from(a_constant) | (u8::from(b_constant) << 1)); + } + + /// Execute `build(a, b)` and assert that `prepare` saw exactly `expect_seen` as its constant + /// operands, so the test knows which decode path the inputs took. + fn run_probed( + build: &impl Fn(ArrayRef, ArrayRef) -> VortexResult, + a: ArrayRef, + b: ArrayRef, + expect_seen: u8, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + SEEN_CONSTANTS.set(u8::MAX); + let result = build(a, b)? + .into_array() + .execute::(ctx)? + .into_array(); + + assert_eq!( + SEEN_CONSTANTS.get(), + expect_seen, + "prepare saw the wrong constant operands", + ); + Ok(result) + } + + /// Assert that every constant-operand arrangement of `build(a, b)` returns exactly what the + /// fully expanded columns return, and that each arrangement's constness really reached + /// `prepare` (so the constants exercised the stride-0 path rather than a decoded column). + /// + /// Arrangements: `a` constant, `b` constant, and both constant with `a` masked. A plain + /// constant pair folds to a single-row execution before the row loop, so masking one side is + /// what drives the both-hoisted arm across rows; that run is compared against the same mask + /// over the expanded column. + pub(crate) fn assert_prepared_agrees_with_columns( + build: impl Fn(ArrayRef, ArrayRef) -> VortexResult, + const_a: ArrayRef, + const_b: ArrayRef, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let col_a = const_a.clone().execute::(&mut ctx)?.into_array(); + let col_b = const_b.clone().execute::(&mut ctx)?.into_array(); + + let baseline = run_probed(&build, col_a.clone(), col_b.clone(), 0b00, &mut ctx)?; + let a_hoisted = run_probed(&build, const_a.clone(), col_b.clone(), 0b01, &mut ctx)?; + let b_hoisted = run_probed(&build, col_a.clone(), const_b.clone(), 0b10, &mut ctx)?; + assert_arrays_eq!(a_hoisted, baseline, &mut ctx); + assert_arrays_eq!(b_hoisted, baseline, &mut ctx); + + let validity = Validity::from_iter((0..col_a.len()).map(|row| row != 1)); + let masked_const_a = MaskedArray::try_new(const_a, validity.clone())?.into_array(); + let masked_col_a = MaskedArray::try_new(col_a, validity)?.into_array(); + let both_hoisted = run_probed(&build, masked_const_a, const_b, 0b11, &mut ctx)?; + let masked_baseline = run_probed(&build, masked_col_a, col_b, 0b00, &mut ctx)?; + assert_arrays_eq!(both_hoisted, masked_baseline, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-tensor/Cargo.toml b/vortex-tensor/Cargo.toml index abdca676775..9706102f6d6 100644 --- a/vortex-tensor/Cargo.toml +++ b/vortex-tensor/Cargo.toml @@ -17,7 +17,7 @@ version = { workspace = true } workspace = true [dependencies] -vortex-array = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } diff --git a/vortex-tensor/benches/cosine_similarity.rs b/vortex-tensor/benches/cosine_similarity.rs index 6cc5eb867ef..49551fbf701 100644 --- a/vortex-tensor/benches/cosine_similarity.rs +++ b/vortex-tensor/benches/cosine_similarity.rs @@ -21,11 +21,14 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity; @@ -41,15 +44,11 @@ fn main() { } /// Total `f64` elements per operand, held constant across widths: the row count is -/// `ELEMENTS / width`. This budget is a quarter of the one the other tensor benches use, because -/// the constant arms recompute the broadcast vector's norm per row and cost roughly ten times the -/// column arms per element. It is what keeps every arm inside the 1 ms per-iteration limit from -/// `docs/developer-guide/benchmarking.md`, measured against CodSpeed's CPU simulation. +/// `ELEMENTS / width`. The smaller budget keeps the wider cosine kernels inside the 1 ms +/// per-iteration limit from `docs/developer-guide/benchmarking.md` under CodSpeed simulation. const ELEMENTS: usize = 2_048; -/// Widths chosen to separate the two costs, as in `l2_norm.rs`: the redundant norm pass is -/// `O(rows * width)`, one third of the closure's arithmetic, so wide tensors show the hoist -/// while a narrow one is dominated by per-row framework costs. +/// Widths that expose both fixed row-framework costs and the `O(width)` kernel work. const WIDTHS: &[usize] = &[2, 32, 256]; /// `ELEMENTS / width` vectors of `width` `f64` elements, non-nullable. `seed` offsets the values so @@ -85,9 +84,9 @@ fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { bencher .with_inputs(|| { ( - CosineSimilarity::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + CosineSimilarity + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -106,6 +105,22 @@ fn column_x_constant(bencher: Bencher, width: usize) { bench_cosine(bencher, vectors(width, 0), constant_vector(width)); } +/// The lhs is a broadcast query vector, whose norm is the same in every row. +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_cosine(bencher, constant_vector(width), vectors(width, 31)); +} + +/// A nullable broadcast rhs exercises constant preparation and output validity together. +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_cosine(bencher, vectors(width, 0), rhs); +} + /// One query vector represented as an extension array over constant storage. fn extension_constant_vector(width: usize) -> ArrayRef { let ext_dtype = vectors(width, 0).dtype().as_extension().clone(); diff --git a/vortex-tensor/benches/inner_product.rs b/vortex-tensor/benches/inner_product.rs index 796e9b648d6..5c4adf1c7ec 100644 --- a/vortex-tensor/benches/inner_product.rs +++ b/vortex-tensor/benches/inner_product.rs @@ -15,11 +15,19 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::inner_product::InnerProduct; @@ -56,15 +64,25 @@ fn vectors(width: usize, seed: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(lhs.len())) .with_inputs(|| { ( - InnerProduct::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + InnerProduct + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -84,3 +102,22 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_inner_product(bencher, lhs, vectors(width, 31)); } + +#[divan::bench(args = WIDTHS)] +fn column_x_constant(bencher: Bencher, width: usize) { + bench_inner_product(bencher, vectors(width, 0), constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_inner_product(bencher, constant_vector(width), vectors(width, 31)); +} + +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_inner_product(bencher, vectors(width, 0), rhs); +} diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index 6f597084113..bf8832f2520 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -15,11 +15,19 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::l2_norm::L2Norm; @@ -54,13 +62,25 @@ fn vectors(width: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(input.len())) .with_inputs(|| { ( - L2Norm::try_new_array(input.clone()).unwrap().into_array(), + L2Norm + .try_new_array(input.len(), EmptyOptions, [input.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -80,3 +100,17 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_l2_norm(bencher, input); } + +#[divan::bench(args = WIDTHS)] +fn constant(bencher: Bencher, width: usize) { + bench_l2_norm(bencher, constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let input = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_l2_norm(bencher, input); +} diff --git a/vortex-tensor/src/encodings/normalized/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs index 0f79ce97bee..ed840a5771f 100644 --- a/vortex-tensor/src/encodings/normalized/mod.rs +++ b/vortex-tensor/src/encodings/normalized/mod.rs @@ -21,7 +21,6 @@ pub use array::NormalizedSlots; mod compress; pub use compress::NormalizedScheme; pub use compress::normalize; -pub(crate) use compress::try_build_constant_normalized; mod execute; diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index ef6eed69e94..1582ef9c7d0 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Cosine similarity expression for tensor-like types. +//! Cosine similarity between two tensor columns. +use num_traits::Float; use num_traits::Zero; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -10,39 +11,43 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; +use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::encodings::normalized::try_build_constant_normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; +use crate::scalar_fns::row::TensorRow; +#[cfg(test)] +use crate::scalar_fns::row::probe; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; +use crate::utils::l2_norm_row; /// Cosine similarity between two columns. /// /// Computes `dot(a, b) / (||a|| * ||b||)` over the flat backing buffer of each tensor or vector. /// The shape and permutation do not affect the result because cosine similarity only depends on the -/// element values, not their logical arrangement. +/// element values, not their logical arrangement. A zero norm on either side yields `0.0`. /// /// Both inputs must be tensor-like extension arrays ([`FixedShapeTensor`] or [`Vector`]) with the /// same dtype and a float element type. The output is a float column of the same float type. @@ -55,13 +60,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct CosineSimilarity; impl CosineSimilarity { /// Creates a new [`TypedScalarFnInstance`] wrapping the cosine similarity operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(CosineSimilarity, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the cosine similarity between `lhs` and @@ -69,127 +74,90 @@ impl CosineSimilarity { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(CosineSimilarity::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for CosineSimilarity { +impl RowFn for CosineSimilarity { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("CosineSimilarity must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow, TensorRow), UninitElementSink, _, _>( + |(lhs, rhs)| { + #[cfg(test)] + probe::record(lhs.is_some(), rhs.is_some()); + ConstantNorms { + lhs: lhs.map(l2_norm_row), + rhs: rhs.map(l2_norm_row), + } + }, + |norms, (lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { + InitializedElement::write( + output, + cosine_similarity_row_prepared(norms, lhs, rhs), + ) + } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands make the _stored_ norms and normalized children + /// authoritative: `cos(D(x, s), D(y, t)) = dot(x, y)` and `cos(D(x, s), y) = dot(x, y) / + /// ||y||`, in both cases forced to `0.0` on rows where any authoritative norm is `0.0` (even + /// for lossy children whose decoded coordinates are nonzero). + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, + args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut lhs_ref = args.get(0)?; - let mut rhs_ref = args.get(1)?; - let len = args.row_count(); - - // Normalize extension-level constants so the encoded fast path can use them. - if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, len, ctx)? { - lhs_ref = normalized_array.into_array(); - } - if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, len, ctx)? { - rhs_ref = normalized_array.into_array(); - } + ) -> VortexResult> { + let lhs = args[0].clone(); + let rhs = args[1].clone(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { - NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); - } + match NormalizedOrientation::classify(&lhs, &rhs) { + NormalizedOrientation::Both { lhs, rhs } => cosine_both_normalized(lhs, rhs, ctx) + .map(|output| Some(RowExecution::Output(output))), NormalizedOrientation::One { normalized_array, plain, - } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); - } - NormalizedOrientation::Neither => {} + } => cosine_one_normalized(normalized_array, plain, ctx) + .map(|output| Some(RowExecution::Output(output))), + NormalizedOrientation::Neither => Ok(None), } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm::try_new_array(lhs_ref.clone())?; - let norm_rhs_arr = L2Norm::try_new_array(rhs_ref.clone())?; - let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; - - // Execute to get the inner product and norms of the arrays. We only fully decompress - // because we need to perform special logic (guard against 0) during division. - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - let norm_l: PrimitiveArray = norm_lhs_arr.into_array().execute(ctx)?; - let norm_r: PrimitiveArray = norm_rhs_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norm_l.as_slice::(); - let norms_r = norm_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - let denom = norms_l[i] * norms_r[i]; - - if denom == T::zero() { - T::zero() - } else { - dots[i] / denom - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false } } @@ -219,584 +187,177 @@ impl ScalarFnArrayVTable for CosineSimilarity { } } -impl CosineSimilarity { - /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so - /// `cosine_similarity = dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - // `Normalized` makes the normalized children authoritative, so their dot product is the - // cosine similarity even for lossy storage wrappers, except that a zero stored norm still - // represents a zero vector. - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norms_l.as_slice::(); - let norms_r = norms_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if norms_l[i] == T::zero() || norms_r[i] == T::zero() { - T::zero() - } else { - dots[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: treat the normalized child as authoritative, so - /// `cosine_similarity = dot(n, b) / ||b||`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, normalized_norms) = extract_normalized_children(normalized_ref); - - let dot_arr = InnerProduct::try_new_array(normalized, plain_ref.clone())?; - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - - let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - - let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; - let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let normalized_norms = normalized_norms.as_slice::(); - let plain_norms = plain_norm.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if normalized_norms[i] == T::zero() || plain_norms[i] == T::zero() { - T::zero() - } else { - dots[i] / plain_norms[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } +/// Per-batch state for the cosine row kernel: the L2 norm of each operand that is constant for +/// the batch. +/// +/// A broadcast query vector holds the same elements in every row, so its norm is the same in +/// every row too. Computing it in the prepare step hoists an `O(width)` pass and a `sqrt` per row +/// out of the row loop. `None` marks an operand that varies by row, whose norm the row closure +/// computes exactly as it did before the hoist. +struct ConstantNorms { + /// The norm of the lhs when it is batch-constant. + lhs: Option, + + /// The norm of the rhs when it is batch-constant. + rhs: Option, } -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::cosine_similarity::CosineSimilarity; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::constant_tensor_array; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. - fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[test] - fn unit_vectors_1d() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 0.0, 1.0, 0.0, // Tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 1.0, 0.0, 0.0, // Tensor 2 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - /// Single-row cosine similarity for various vector pairs. - #[rstest] - // Antiparallel -> -1.0. - #[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] - // dot=24, both magnitudes=5 -> 24/25 = 0.96. - #[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] - // Zero vector -> guarded to 0.0. - #[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); - Ok(()) - } - - /// Self-similarity across various tensor shapes should always produce 1.0. - #[rstest] - // 2x3 matrix, flattened to 6 elements. - #[case::matrix_2d( - &[2, 3], - &[ - 1.0, 0.0, 0.0, // row 0 - 0.0, 0.0, 0.0, // row 1 - ], - )] - // 2x2x2 tensor, 8 elements. - #[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] - fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { - let lhs = tensor_array(shape, elements)?; - let rhs = tensor_array(shape, elements)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn scalar_0d() -> VortexResult<()> { - // 0-dimensional tensor: each "tensor" is a single scalar value. - let lhs = tensor_array(&[], &[5.0, 3.0])?; - let rhs = tensor_array(&[], &[5.0, -3.0])?; - - // Same sign -> 1.0, opposite sign -> -1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); - Ok(()) - } - - #[test] - fn many_rows() -> VortexResult<()> { - // 5 tensors of shape [4] compared against themselves -> all 1.0. - let lhs = tensor_array( - &[4], - &[ - 1.0, 2.0, 3.0, 4.0, // tensor 0 - 0.0, 1.0, 0.0, 0.0, // tensor 1 - 5.0, 0.0, 5.0, 0.0, // tensor 2 - 1.0, 1.0, 1.0, 1.0, // tensor 3 - 0.0, 0.0, 0.0, 7.0, // tensor 4 - ], - )?; - let rhs = lhs.clone(); - - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0, 1.0, 1.0, 1.0, 1.0], - ); - Ok(()) - } - - #[test] - fn constant_query_tensor() -> VortexResult<()> { - // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. - let data = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 0.0, 1.0, 0.0, // tensor 1 - 0.0, 0.0, 1.0, // tensor 2 - 1.0, 0.0, 0.0, // tensor 3 - ], - )?; - let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn vector_unit_vectors() -> VortexResult<()> { - let lhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 1.0, 0.0, 0.0, // vector 1 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn vector_constant_query() -> VortexResult<()> { - let data = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - 0.0, 0.0, 1.0, // vector 2 - 1.0, 0.0, 0.0, // vector 3 - ], - )?; - let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. - let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; - let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; - let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: self-similarity = 1.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_self_similarity() -> VortexResult<()> { - // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. - // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Self-similarity should always be 1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); - Ok(()) - } - - #[test] - fn both_normalized_orthogonal() -> VortexResult<()> { - // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. - // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn both_normalized_zero_norm() -> VortexResult<()> { - // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS is plain [3.0, 4.0]. - // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[3.0, 4.0])?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 0.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); - Ok(()) - } - - #[test] - fn both_normalized_null_rows() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_r = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - let validity = Validity::from_iter([true, false]); - let rhs = Normalized::try_new(normalized_r, norms_r, validity, &mut ctx)?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine - // similarity for that row must be `0.0` even though the dot product of the normalized - // children is nonzero. - let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; - let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally violates the unit-norm invariant by pairing a nonzero normalized row - // with a stored norm of `0.0`, mimicking lossy storage. - // SAFETY: The children are structurally valid. - let lhs = - unsafe { Normalized::new_unchecked(normalized_l, norms_l, Validity::NonNullable) } - .into_array(); - - let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; - let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Same as above for the rhs operand. - // SAFETY: The children are structurally valid. - let rhs = - unsafe { Normalized::new_unchecked(normalized_r, norms_r, Validity::NonNullable) } - .into_array(); - - // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both - // `0.0`, so cosine similarity must be `0.0`. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. The plain side is a normal nonzero - // tensor with positive norm. cosine similarity must still be `0.0` because the - // authoritative stored norm on the normalized_array side is `0.0`. - let normalized = tensor_array(&[2], &[0.6, 0.8])?; - let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally pairs a nonzero normalized row with a stored norm of `0.0`, mimicking - // lossy storage where the stored norm is authoritative. - // SAFETY: The children are structurally valid. - let normalized_array = - unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) } - .into_array(); - - let plain = tensor_array(&[2], &[1.0, 0.0])?; - - // Normalized encoding on the lhs: `One { normalized_array: lhs, plain: rhs }`. - assert_close( - &eval_cosine_similarity(normalized_array.clone(), plain.clone())?, - &[0.0], - ); - - // Normalized encoding on the rhs: `One { normalized_array: rhs, plain: lhs }`. The same - // zero-norm guard must fire regardless of operand order. - assert_close(&eval_cosine_similarity(plain, normalized_array)?, &[0.0]); - Ok(()) - } - - #[test] - fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { - // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. - // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. - let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 - 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 - 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 - 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { - // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn both_constant_tensors() -> VortexResult<()> { - // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). - let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; - let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; - let expected = 1.0 / 2.0_f64.sqrt(); - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[expected, expected, expected], - ); - Ok(()) - } - - #[test] - fn constant_zero_norm_query() -> VortexResult<()> { - // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in - // `execute_one_normalized` and `execute_both_normalized`. - let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 2.0, 3.0, // - 4.0, 5.0, 6.0, // - 7.0, 8.0, 9.0, // - ], - )?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); - Ok(()) - } - - #[test] - fn constant_self_similarity_nonunit() -> VortexResult<()> { - // A non-unit constant query compared to itself must produce `1.0`. This exercises the - // helper's division: after normalization, both sides must be exactly unit so the - // Normalized fast path's inner product yields 1. - let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); - Ok(()) - } - - #[test] - fn vector_constant_matches_plain() -> VortexResult<()> { - // Exercise the `Vector` extension variant through the new pre-pass. - let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[rstest] - #[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] - #[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = CosineSimilarity::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("CosineSimilarity serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) +/// Computes the cosine similarity of one row, taking any hoisted norm from `norms` and computing +/// the rest exactly as [`cosine_similarity_row`] does. +/// +/// Each arm accumulates the same values in the same order as [`cosine_similarity_row`], and the +/// denominator keeps its lhs-times-rhs order, so the result is bit-identical whether a norm was +/// hoisted or not. The match costs one predictable branch per row: the arm is the same for the +/// whole batch. +fn cosine_similarity_row_prepared( + norms: &ConstantNorms, + lhs: &[T], + rhs: &[T], +) -> T { + match (norms.lhs, norms.rhs) { + (None, None) => cosine_similarity_row(lhs, rhs), + (Some(lhs_norm), None) => { + let mut dot = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm_squared.sqrt()) + } + (None, Some(rhs_norm)) => { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + } + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm) + } + (Some(lhs_norm), Some(rhs_norm)) => { + let mut dot = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm) + } } +} - fn cosine_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") - } +/// Computes the cosine similarity of two equal-length float slices. +/// +/// Returns `dot(a, b) / (||a|| * ||b||)`, or `0.0` when either norm is zero. +fn cosine_similarity_row(lhs: &[T], rhs: &[T]) -> T { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm_squared.sqrt()) +} - fn cosine_vector_rhs() -> ArrayRef { - vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +/// The shared tail of every cosine arm: `dot / denominator`, guarded to `0.0` when it is +/// zero. +fn cosine_from_parts(dot: T, denominator: T) -> T { + if denominator == T::zero() { + T::zero() + } else { + dot / denominator } +} - fn cosine_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") - } +/// Both sides are [`Normalized`]-encoded: the normalized children are authoritative, so their dot +/// product is the cosine similarity, except that a row with a zero _stored_ norm is a zero vector. +/// +/// Unlike [`InnerProduct::reduce_encoded`], which composes lazy `Mul` arrays over the norm columns, +/// this executes and materializes. The zero-norm guard is a conditional per row rather than an +/// arithmetic factor, so there is no lazy array that expresses it; the norm columns are one value +/// per row rather than one per coordinate, so materializing them is cheap next to the decode this +/// avoids. +/// +/// [`InnerProduct::reduce_encoded`]: InnerProduct::reduce_encoded +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_both_normalized( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized_l, normalized_r])? + .execute(ctx)?; + let norms_l: PrimitiveArray = norms_l.execute(ctx)?; + let norms_r: PrimitiveArray = norms_r.execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let norms_l = norms_l.as_slice::(); + let norms_r = norms_r.as_slice::(); + // Zipped rather than indexed by `0..len`: one bounds check per iterator instead of three + // per row. A length disagreement between the children shortens the result, which the + // lifting reports against the batch row count rather than panicking mid-loop. + let buffer: Buffer = dots + .iter() + .zip(norms_l) + .zip(norms_r) + .map(|((&dot, &norm_l), &norm_r)| { + if norm_l.is_zero() || norm_r.is_zero() { + T::zero() + } else { + dot + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) +} - fn cosine_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") - } +/// One side is [`Normalized`]-encoded: `cos = dot(normalized, plain) / ||plain||`, forced to `0.0` +/// on rows where the stored norm or the plain norm is `0.0`. +/// +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_one_normalized( + normalized_array: &ArrayRef, + plain: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = normalized_array.len(); + let (normalized, normalized_norms) = extract_normalized_children(normalized_array); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized, plain.clone()])? + .execute(ctx)?; + let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; + let plain_norm: PrimitiveArray = L2Norm + .try_new_array(len, EmptyOptions, [plain.clone()])? + .execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let normalized_norms = normalized_norms.as_slice::(); + let plain_norms = plain_norm.as_slice::(); + // Zipped for the same reason as [`cosine_both_normalized`]. + let buffer: Buffer = dots + .iter() + .zip(normalized_norms) + .zip(plain_norms) + .map(|((&dot, &stored_norm), &plain_norm)| { + if stored_norm.is_zero() || plain_norm.is_zero() { + T::zero() + } else { + dot / plain_norm + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) } diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 74eb184045f..a31cce5ddda 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -6,40 +6,34 @@ use num_traits::Float; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::matcher::AnyTensor; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; -use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; /// Inner product (dot product) between two columns. /// @@ -52,13 +46,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct InnerProduct; impl InnerProduct { /// Creates a new [`TypedScalarFnInstance`] wrapping the inner product operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(InnerProduct, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the inner product between `lhs` and @@ -66,116 +60,89 @@ impl InnerProduct { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(InnerProduct::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for InnerProduct { +impl RowFn for InnerProduct { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.inner_product"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("InnerProduct must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - // TODO(connor): relax the float-only gate once integer tensors are supported. - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow, TensorRow), UninitElementSink, _>( + |(lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, inner_product_row(lhs, rhs)) } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands factor through their stored norms: with `D(x, s)` denoting + /// `x * s` rowwise, `dot(D(x, s), D(y, t)) = s * t * dot(x, y)` and + /// `dot(D(x, s), y) = s * dot(x, y)`. The rewrite is expressed with lazy [`Operator::Mul`] + /// arrays over the (much smaller) norm columns, so no denormalized coordinates are decoded. + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let lhs_ref = args.get(0)?; - let rhs_ref = args.get(1)?; - let len = args.row_count(); + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let len = args[0].len(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + Ok(match NormalizedOrientation::classify(&args[0], &args[1]) { NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized_l, normalized_r])?; + Some( + dot.binary(norms_l, Operator::Mul)? + .binary(norms_r, Operator::Mul)?, + ) } NormalizedOrientation::One { normalized_array, plain, } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); + let (normalized, norms) = extract_normalized_children(normalized_array); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized, plain.clone()])?; + Some(dot.binary(norms, Operator::Mul)?) } - NormalizedOrientation::Neither => {} + NormalizedOrientation::Neither => None, } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Canonicalize so we can perform the math directly. - let lhs: ExtensionArray = lhs_ref.execute(ctx)?; - let rhs: ExtensionArray = rhs_ref.execute(ctx)?; - - // We validated that both inputs have the same type. - let ext = lhs.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let dimensions = tensor_match.list_size() as usize; - - // Extract the storage array from each extension input. We pass the storage (FSL) rather - // than the extension array to avoid canonicalizing the extension wrapper. - let lhs_storage = lhs.storage_array(); - let rhs_storage = rhs.storage_array(); - - let lhs_flat = extract_flat_elements(lhs_storage, dimensions, ctx)?; - let rhs_flat = extract_flat_elements(rhs_storage, dimensions, ctx)?; - - match_each_float_ptype!(lhs_flat.ptype(), |T| { - let buffer: Buffer = (0..len) - .map(|i| inner_product_row(lhs_flat.row::(i), rhs_flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + .map(RowExecution::Output)) } } @@ -205,329 +172,12 @@ impl ScalarFnArrayVTable for InnerProduct { } } -impl InnerProduct { - /// Both sides are [`Normalized`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let nl = norms_l.as_slice::(); - let nr = norms_r.as_slice::(); - let buffer: Buffer = (0..len).map(|i| nl[i] * nr[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: `inner_product = s * dot(n, other)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, norms) = extract_normalized_children(normalized_ref); - let normalized_norms: PrimitiveArray = norms.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized, plain_ref.clone())? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let ns = normalized_norms.as_slice::(); - let buffer: Buffer = (0..len).map(|i| ns[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } -} - /// Computes the inner product (dot product) of two equal-length float slices. /// /// Returns `sum(a_i * b_i)`. -fn inner_product_row(a: &[T], b: &[T]) -> T { - a.iter() - .zip(b.iter()) - .map(|(&x, &y)| x * y) - .fold(T::zero(), |acc, v| acc + v) -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::inner_product::InnerProduct; - use crate::tests::SESSION; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. - fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - /// Single-row inner product for various vector pairs. - #[rstest] - // Orthogonal: [1, 0] . [0, 1] = 0. - #[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] - // Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. - #[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] - // Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. - #[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] - // Scaled: [2, 0] . [3, 0] = 6. - #[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_inner_product(lhs, rhs)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 3.0, 4.0, 0.0, // tensor 1 - 1.0, 1.0, 1.0, // tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 0.0, 1.0, 0.0, // tensor 0: dot = 0 - 3.0, 4.0, 0.0, // tensor 1: dot = 25 - 2.0, 2.0, 2.0, // tensor 2: dot = 6 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); - Ok(()) - } - - #[test] - fn vector_inner_product() -> VortexResult<()> { - let lhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0 - 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0: dot = 25 - 0.0, 1.0, // vector 1: dot = 0 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. - let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; - let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; - let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert!(prim.is_valid(2, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[23.0]); - assert_close(&[prim.as_slice::()[2]], &[127.0]); - Ok(()) - } - - #[test] - fn rejects_non_extension_dtype() { - let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); - let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - } - - #[test] - fn rejects_mismatched_dtypes() -> VortexResult<()> { - let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; - let rhs = vector_array(2, &[3.0_f64, 4.0])?; - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn both_normalized() -> VortexResult<()> { - // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). - // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). - // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; - - // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. - assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); - Ok(()) - } - - #[test] - fn both_normalized_multiple_rows() -> VortexResult<()> { - // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. - // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS: plain [1.0, 2.0]. - // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[1.0, 2.0])?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS: plain [1.0, 2.0]. - // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 2.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn both_normalized_null_rows() -> VortexResult<()> { - let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_l = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let validity = Validity::from_iter([true, false]); - let lhs = Normalized::try_new(normalized_l, norms_l, validity, &mut ctx)?.into_array(); - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[25.0]); - Ok(()) - } - - #[rstest] - #[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] - #[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = InnerProduct::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(InnerProduct); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("InnerProduct serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn inner_product_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } - - fn inner_product_vector_rhs() -> ArrayRef { - vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") - } - - fn inner_product_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") - } - - fn inner_product_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") - } +fn inner_product_row(lhs: &[T], rhs: &[T]) -> T { + lhs.iter() + .zip(rhs) + .map(|(&lhs_element, &rhs_element)| lhs_element * rhs_element) + .fold(T::zero(), |sum, product| sum + product) } diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index a72b9da66db..78cabc961b0 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -3,51 +3,44 @@ //! L2 norm expression for tensor-like types. -use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; -use crate::matcher::AnyTensor; -use crate::utils::extract_flat_elements; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::extract_normalized_children; -use crate::utils::reattach_validity; +use crate::utils::l2_norm_row; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -63,139 +56,112 @@ use crate::utils::validate_tensor_float_input; /// of the storage contract, not a separate lossy-compute mode. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct L2Norm; impl L2Norm { /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 norm operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(L2Norm, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the L2 norm over `child`. /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtype is + /// unsupported. pub fn try_new_array(child: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(L2Norm::new().erased(), vec![child]) + ScalarFnArray::try_new(Self::new().erased(), vec![child]) } } -impl ScalarFnVTable for L2Norm { +impl RowFn for L2Norm { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.l2_norm"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(1) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("input"), - _ => unreachable!("L2Norm must have exactly one child"), - } - } - - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let input_dtype = &arg_dtypes[0]; - let tensor_match = validate_tensor_float_input(input_dtype)?; - let ptype = tensor_match.element_ptype(); - - let nullability = Nullability::from(input_dtype.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn execute( + fn dispatch>( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let input_ref = args.get(0)?; - let row_count = args.row_count(); - - let ext = input_ref.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let tensor_flat_size = tensor_match.list_size() as usize; - let element_ptype = tensor_match.element_ptype(); - - let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); - - // Stored norms are authoritative. Reattach the parent validity because the child is - // non-nullable. - if input_ref.is::() { - let (_, norms) = extract_normalized_children(&input_ref); - let norms = reattach_validity(norms, input_ref.validity()?)?; - vortex_ensure_eq!(norms.dtype(), &norm_dtype); - return Ok(norms); - } - - // Optimize for the constant array case. - if let Some(array) = input_ref.as_opt::() { - let scalar = array.scalar().as_extension().to_storage_scalar(); - - let Some(elements) = scalar.as_list().elements() else { - return Ok(ConstantArray::new(Scalar::null(norm_dtype), row_count).into_array()); - }; - - let norm_scalar = match_each_float_ptype!(element_ptype, |T| { - let values: Vec = elements - .iter() - .map(|s| { - s.as_primitive() - .as_::() - .vortex_expect("element was somehow not the correct float") - }) - .collect(); - let norm = l2_norm_row::(&values); - - Scalar::try_new(norm_dtype, Some(norm.into())) - })?; - - let norms = ConstantArray::new(norm_scalar, row_count).into_array(); - return Ok(norms); - } - - let input: ExtensionArray = input_ref.execute(ctx)?; - let validity = input.as_ref().validity()?; - - let storage = input.storage_array(); - let flat = extract_flat_elements(storage, tensor_flat_size, ctx)?; - - match_each_float_ptype!(flat.ptype(), |T| { - let buffer: Buffer = (0..row_count) - .map(|i| l2_norm_row(flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l2_norm_row(row)) } + }) }) } - fn validity( + /// `L2Norm` over a [`Normalized`]-encoded column is defined to read back the authoritative + /// stored norms. Callers of lossy encodings opt into that storage semantics instead of forcing + /// a decode-and-recompute path here. + fn reduce_encoded( &self, _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if the input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let input = &args[0]; + if input.is::() { + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let (_, norms) = extract_normalized_children(input); + vortex_ensure!( + norms.dtype().is_primitive(), + "normalized norms must be primitive, got {}", + norms.dtype(), + ); + vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); + return Ok(Some(RowExecution::Output(norms))); + } - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + let Some(constant) = input.as_opt::() else { + return Ok(None); + }; + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let norm_dtype = + DType::Primitive(element_ptype, input.dtype().as_extension().nullability()); + let storage = constant.scalar().as_extension().to_storage_scalar(); + + let Some(elements) = storage.as_list().elements() else { + let output = ConstantArray::new(Scalar::null(norm_dtype), input.len()); + return Ok(Some(RowExecution::Output(output.into_array()))); + }; + + let norm = match_each_float_ptype!(element_ptype, |T| { + let values: Vec = elements + .iter() + .map(|element| { + element + .as_primitive() + .as_::() + .vortex_expect("tensor element must match its declared ptype") + }) + .collect(); + Scalar::try_new(norm_dtype, Some(l2_norm_row::(&values).into())) + })?; + let output = ConstantArray::new(norm, input.len()); + Ok(Some(RowExecution::Output(output.into_array()))) } } @@ -241,230 +207,3 @@ impl ScalarFnArrayVTable for L2Norm { }) } } - -/// Computes the L2 norm (Euclidean norm) of a float slice. -/// -/// Returns `sqrt(sum(v_i^2))`. A zero-length or all-zero input produces `0.0`. -fn l2_norm_row(v: &[T]) -> T { - let mut sum_sq = T::zero(); - for &x in v { - sum_sq = sum_sq + x * x; - } - sum_sq.sqrt() -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::EmptyMetadata; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Constant; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::dtype::extension::ExtDType; - use vortex_array::scalar::Scalar; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::l2_norm::L2Norm; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::literal_vector_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. - fn eval_l2_norm(input: ArrayRef) -> VortexResult> { - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[rstest] - #[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] - #[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] - #[case::single_element(&[1], &[7.0], &[7.0])] - #[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] - fn known_norms( - #[case] shape: &[usize], - #[case] elements: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let arr = tensor_array(shape, elements)?; - assert_close(&eval_l2_norm(arr)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let arr = tensor_array( - &[3], - &[ - 3.0, 4.0, 0.0, // norm = 5.0 - 0.0, 0.0, 0.0, // norm = 0.0 - 1.0, 1.0, 1.0, // norm = sqrt(3) - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); - Ok(()) - } - - #[test] - fn vector_multiple_rows() -> VortexResult<()> { - let arr = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // norm = 1.0 - 3.0, 4.0, 0.0, // norm = 5.0 - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 is masked as null. - let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; - let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: norm = 5.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is a non-null tensor should short-circuit to a - /// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so - /// execution stops at the [`Constant`] encoding instead of canonicalizing into a - /// [`PrimitiveArray`]. - #[test] - fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { - let input = literal_vector_array(&[3.0f64, 4.0], 4); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("L2Norm over a constant input must produce a constant output"); - assert_eq!(constant.len(), 4); - let norm = constant - .scalar() - .as_primitive() - .as_::() - .expect("norm scalar must be a non-null primitive"); - assert_close(&[norm], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of - /// the correct primitive dtype and length. - #[test] - fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { - let storage_dtype = DType::FixedSizeList( - DType::Primitive(PType::F64, Nullability::NonNullable).into(), - 2, - Nullability::Nullable, - ); - let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); - let null_scalar = Scalar::null(DType::Extension(ext_dtype)); - let input = ConstantArray::new(null_scalar, 3).into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("null constant input must produce a constant output"); - assert_eq!(constant.len(), 3); - assert!(constant.scalar().is_null()); - assert_eq!( - constant.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - Ok(()) - } - - #[test] - fn reads_through_a_nullable_normalized_column() -> VortexResult<()> { - let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; - let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - let validity = Validity::from_iter([true, false]); - let input = Normalized::try_new(normalized, norms, validity, &mut ctx)?.into_array(); - - let result = ScalarFnArray::try_new(L2Norm::new().erased(), vec![input])?.into_array(); - let prim: PrimitiveArray = result.execute(&mut ctx)?; - - assert_eq!( - prim.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - - Ok(()) - } - - #[rstest] - #[case::fixed_shape_tensor(l2_norm_tensor_child())] - #[case::vector(l2_norm_vector_child())] - fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { - let original = L2Norm::try_new_array(child.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(L2Norm); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("L2Norm serialize must produce metadata"); - - let children = vec![child]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn l2_norm_tensor_child() -> ArrayRef { - tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") - } - - fn l2_norm_vector_child() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } -} diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 68f10ca6b01..da9b8950e7a 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -6,3 +6,7 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; +pub(crate) mod row; + +#[cfg(test)] +mod tests; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs new file mode 100644 index 00000000000..0b3cf3634ee --- /dev/null +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the tensor scalar functions add to the row-function machinery: an element type that reads a +//! tensor row and the width rule they share. + +use std::marker::PhantomData; + +use num_traits::Float; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::Masked; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::masked::MaskedArraySlotsExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::scalar_fn::unstable::row::InputElement; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; +use crate::utils::validate_tensor_float_inputs; + +/// The width rule the tensor scalar functions share: every argument is the same float tensor dtype, +/// and the width is its element ptype. +pub(crate) fn tensor_element_ptype(args: &[DType]) -> VortexResult { + Ok(validate_tensor_float_inputs(args)?.element_ptype()) +} + +/// Marker for tensor-valued input elements: accepts any tensor-like extension column whose +/// elements are `T`, and presents each row as its flat elements, `&[T]`. +pub struct TensorRow(PhantomData); + +/// The decoded form of a [`TensorRow`] column: one flat typed buffer plus the stride to read it at. +/// +/// Typed at decode time rather than per row. `FlatElements::row` re-derives its typed slice on every +/// call, which costs a ptype check and a buffer downcast per row; a row loop reads every row, so it +/// pays that once here instead. +pub struct TensorRows { + /// Every row's elements, back to back. + elements: Buffer, + + /// Number of logical tensor rows, stored so zero-width tensors retain their length. + rows: usize, + + /// Elements per row, the length of each row slice. + list_size: usize, + + /// `list_size` for a full column and `0` for constant-backed storage, so `index * stride` pins a + /// constant to its single materialized row without a branch in the loop. + stride: usize, +} + +// SAFETY: `TensorRows` records the row count validated during decode, and both checked and +// unchecked access use the same stride and row width. +unsafe impl InputElement for TensorRow { + type Column = TensorRows; + type View<'a> = &'a TensorRows; + type Elem<'a> = &'a [T]; + + // Tensor storage is a fully materialized non-nullable primitive buffer, so the elements behind + // a null row are arbitrary values rather than an unresolvable reference. + const DENSE_SAFE: bool = true; + // Tensor storage is a primitive buffer; reading it cannot fail on account of its values. + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let tensor_match = validate_tensor_float_input(dtype)?; + let expected = T::PTYPE; + vortex_ensure_eq!( + tensor_match.element_ptype(), + expected, + "expected a tensor of {expected} elements, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Dense batch execution owns the mask and restores it on the result. Decode the values + // directly so a nullable tensor does not rebuild its extension storage under that mask. + let array = match array.as_opt::() { + Some(masked) => masked.child().clone(), + None => array, + }; + + let rows = array.len(); + let list_size = validate_tensor_float_input(array.dtype())?.list_size() as usize; + let ext: ExtensionArray = array.execute(ctx)?; + let flat = extract_flat_elements(ext.storage_array(), list_size, ctx)?; + let list_size = flat.list_size(); + let stride = flat.row_stride(); + let elements = flat.into_buffer::(); + + let expected_elements = if stride == 0 { + list_size + } else { + vortex_ensure_eq!( + stride, + list_size, + "per-row tensor stride must equal its width, got {stride}", + ); + let Some(expected_elements) = rows.checked_mul(stride) else { + vortex_bail!( + "tensor row storage length must fit usize, got {rows} rows of width {stride}", + ); + }; + + expected_elements + }; + vortex_ensure_eq!( + elements.len(), + expected_elements, + "tensor row storage must contain {expected_elements} elements, got {}", + elements.len(), + ); + + Ok(TensorRows { + elements, + rows, + list_size, + stride, + }) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> &[T] { + let start = index * column.stride; + &column.elements.as_slice()[start..start + column.list_size] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.rows + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + Self::get(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + let start = index * view.stride; + + // SAFETY: decode established one complete stored row for stride 0, or `rows` contiguous + // `list_size`-element rows otherwise. The caller guarantees `index < rows`. + unsafe { + std::slice::from_raw_parts(view.elements.as_slice().as_ptr().add(start), view.list_size) + } + } +} + +/// Records which operands a test's prepare step received as batch constants. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + thread_local! { + /// Bit 0 records the lhs and bit 1 records the rhs. Thread-local storage prevents + /// concurrent tests from racing; row execution remains on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Records whether each operand was constant for the current batch. + pub(crate) fn record(lhs_constant: bool, rhs_constant: bool) { + SEEN_CONSTANTS.set(u8::from(lhs_constant) | (u8::from(rhs_constant) << 1)); + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs new file mode 100644 index 00000000000..13eed7b0224 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -0,0 +1,610 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::assert_arrays_eq; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::cosine_similarity::CosineSimilarity; +use crate::scalar_fns::row::probe; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::constant_tensor_array; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. +fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = CosineSimilarity::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = CosineSimilarity::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_cosine_similarity(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Like [`eval_cosine_similarity`], but returns the executed array for exact array comparisons. +fn eval_cosine_similarity_array( + lhs: ArrayRef, + rhs: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + Ok(result + .into_array() + .execute::(ctx)? + .into_array()) +} + +#[test] +fn unit_vectors_1d() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 0.0, 1.0, 0.0, // Tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 1.0, 0.0, 0.0, // Tensor 2 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +/// Single-row cosine similarity for various vector pairs. +#[rstest] +// Antiparallel -> -1.0. +#[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] +// dot=24, both magnitudes=5 -> 24/25 = 0.96. +#[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] +// Zero vector -> guarded to 0.0. +#[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); + Ok(()) +} + +/// Self-similarity across various tensor shapes should always produce 1.0. +#[rstest] +// 2x3 matrix, flattened to 6 elements. +#[case::matrix_2d( + &[2, 3], + &[ + 1.0, 0.0, 0.0, // row 0 + 0.0, 0.0, 0.0, // row 1 + ], +)] +// 2x2x2 tensor, 8 elements. +#[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] +fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { + let lhs = tensor_array(shape, elements)?; + let rhs = tensor_array(shape, elements)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn scalar_0d() -> VortexResult<()> { + // 0-dimensional tensor: each "tensor" is a single scalar value. + let lhs = tensor_array(&[], &[5.0, 3.0])?; + let rhs = tensor_array(&[], &[5.0, -3.0])?; + + // Same sign -> 1.0, opposite sign -> -1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); + Ok(()) +} + +#[test] +fn many_rows() -> VortexResult<()> { + // 5 tensors of shape [4] compared against themselves -> all 1.0. + let lhs = tensor_array( + &[4], + &[ + 1.0, 2.0, 3.0, 4.0, // tensor 0 + 0.0, 1.0, 0.0, 0.0, // tensor 1 + 5.0, 0.0, 5.0, 0.0, // tensor 2 + 1.0, 1.0, 1.0, 1.0, // tensor 3 + 0.0, 0.0, 0.0, 7.0, // tensor 4 + ], + )?; + let rhs = lhs.clone(); + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0, 1.0, 1.0, 1.0, 1.0], + ); + Ok(()) +} + +#[test] +fn constant_query_tensor() -> VortexResult<()> { + // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. + let data = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 0.0, 1.0, 0.0, // tensor 1 + 0.0, 0.0, 1.0, // tensor 2 + 1.0, 0.0, 0.0, // tensor 3 + ], + )?; + let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn vector_unit_vectors() -> VortexResult<()> { + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 1.0, 0.0, 0.0, // vector 1 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn vector_constant_query() -> VortexResult<()> { + let data = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + 0.0, 0.0, 1.0, // vector 2 + 1.0, 0.0, 0.0, // vector 3 + ], + )?; + let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. + let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; + let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; + let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: self-similarity = 1.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_self_similarity() -> VortexResult<()> { + // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. + // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Self-similarity should always be 1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); + Ok(()) +} + +#[test] +fn both_normalized_orthogonal() -> VortexResult<()> { + // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. + // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn both_normalized_zero_norm() -> VortexResult<()> { + // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS is plain [3.0, 4.0]. + // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[3.0, 4.0])?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 0.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); + Ok(()) +} + +#[test] +fn both_normalized_null_rows() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_r = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let validity = Validity::from_iter([true, false]); + let rhs = Normalized::try_new(normalized_r, norms_r, validity, &mut ctx)?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine + // similarity for that row must be `0.0` even though the dot product of the normalized + // children is nonzero. + let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; + let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally violates the unit-norm invariant by + // pairing a nonzero normalized row with a stored norm of `0.0`, mimicking lossy storage. + let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l, Validity::NonNullable) } + .into_array(); + + let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; + let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: Same as above for the rhs operand. + let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r, Validity::NonNullable) } + .into_array(); + + // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both + // `0.0`, so cosine similarity must be `0.0`. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. The plain side is a normal nonzero + // tensor with positive norm. cosine similarity must still be `0.0` because the + // authoritative stored norm on the denorm side is `0.0`. + let normalized = tensor_array(&[2], &[0.6, 0.8])?; + let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally pairs a nonzero normalized row with a + // stored norm of `0.0`, mimicking lossy storage where the stored norm is authoritative. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) }.into_array(); + + let plain = tensor_array(&[2], &[1.0, 0.0])?; + + // Denorm on the lhs: `One { denorm: lhs, plain: rhs }`. + assert_close( + &eval_cosine_similarity(denorm.clone(), plain.clone())?, + &[0.0], + ); + + // Denorm on the rhs: `One { denorm: rhs, plain: lhs }`. The same zero-norm guard must + // fire regardless of operand order. + assert_close(&eval_cosine_similarity(plain, denorm)?, &[0.0]); + Ok(()) +} + +#[test] +fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { + // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. + // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. + let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 + 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 + 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 + 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 + ], + )?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { + // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn both_constant_tensors() -> VortexResult<()> { + // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). + let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; + let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[expected, expected, expected], + ); + Ok(()) +} + +#[test] +fn constant_zero_norm_query() -> VortexResult<()> { + // A zero-norm constant query must produce `0.0` through the prepared row kernel's + // zero-denominator guard. + let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 2.0, 3.0, // + 4.0, 5.0, 6.0, // + 7.0, 8.0, 9.0, // + ], + )?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[test] +fn constant_self_similarity_nonunit() -> VortexResult<()> { + // The prepared path hoists both norms and computes the same dot product for every row. + let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); + Ok(()) +} + +/// An extension array over constant storage (what [`Vector::constant_array`] builds) is a batch +/// constant like any other. The row layer sees through the wrapper, so `prepare` hoists its norm. +#[test] +fn vector_constant_matches_plain() -> VortexResult<()> { + let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + assert_eq!( + probe::SEEN_CONSTANTS.get(), + 0b01, + "the extension-over-constant lhs must reach prepare as a batch constant", + ); + Ok(()) +} + +/// Both literal and extension-wrapped constant storage reach the prepared row path. The probe +/// ensures that the literal query remains a batch constant instead of becoming a per-row column. +/// +/// [`ConstantArray`]: vortex_array::arrays::ConstantArray +#[test] +fn literal_constant_rhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(lhs.clone(), literal_vector_array(&query, 4), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(lhs, vector_array(3, &query.repeat(4))?, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// The mirror of [`literal_constant_rhs_matches_expanded_column`], exercising the hoisted-lhs arm +/// of the prepared kernel. +#[test] +fn literal_constant_lhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(literal_vector_array(&query, 4), rhs.clone(), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(vector_array(3, &query.repeat(4))?, rhs, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// A zero-norm literal constant query must be guarded to `0.0` on every row by the prepared row +/// kernel, exactly as the unprepared kernel guards it. +#[test] +fn literal_constant_zero_norm_query_yields_zero() -> VortexResult<()> { + let lhs = vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = literal_vector_array(&[0.0f64, 0.0, 0.0], 2); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0]); + Ok(()) +} + +/// Two literal constants are folded to a single-row execution by the row lifting, and that row +/// still runs the prepared kernel with both norms hoisted. +#[test] +fn both_literal_constants() -> VortexResult<()> { + let lhs = literal_vector_array(&[1.0f64, 0.0, 0.0], 3); + let rhs = literal_vector_array(&[1.0f64, 1.0, 0.0], 3); + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[expected; 3]); + Ok(()) +} + +#[rstest] +#[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] +#[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("CosineSimilarity serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn cosine_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_vector_rhs() -> ArrayRef { + vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn cosine_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/inner_product.rs b/vortex-tensor/src/scalar_fns/tests/inner_product.rs new file mode 100644 index 00000000000..03d23aaacad --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/inner_product.rs @@ -0,0 +1,282 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::inner_product::InnerProduct; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates inner product between two tensor arrays and returns the result as `Vec`. +fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = InnerProduct::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = InnerProduct::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_inner_product(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Single-row inner product for various vector pairs. +#[rstest] +// Orthogonal: [1, 0] . [0, 1] = 0. +#[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] +// Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. +#[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] +// Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. +#[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] +// Scaled: [2, 0] . [3, 0] = 6. +#[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_inner_product(lhs, rhs)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 3.0, 4.0, 0.0, // tensor 1 + 1.0, 1.0, 1.0, // tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 0.0, 1.0, 0.0, // tensor 0: dot = 0 + 3.0, 4.0, 0.0, // tensor 1: dot = 25 + 2.0, 2.0, 2.0, // tensor 2: dot = 6 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); + Ok(()) +} + +#[test] +fn vector_inner_product() -> VortexResult<()> { + let lhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0 + 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0: dot = 25 + 0.0, 1.0, // vector 1: dot = 0 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. + let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; + let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert!(prim.is_valid(2, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[23.0]); + assert_close(&[prim.as_slice::()[2]], &[127.0]); + Ok(()) +} + +#[test] +fn rejects_non_extension_dtype() { + let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); + let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); +} + +#[test] +fn rejects_mismatched_dtypes() -> VortexResult<()> { + let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; + let rhs = vector_array(2, &[3.0_f64, 4.0])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn both_normalized() -> VortexResult<()> { + // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). + // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). + // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; + + // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. + assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); + Ok(()) +} + +#[test] +fn both_normalized_multiple_rows() -> VortexResult<()> { + // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. + // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS: plain [1.0, 2.0]. + // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[1.0, 2.0])?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS: plain [1.0, 2.0]. + // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 2.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn both_normalized_null_rows() -> VortexResult<()> { + let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_l = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let validity = Validity::from_iter([true, false]); + let lhs = Normalized::try_new(normalized_l, norms_l, validity, &mut ctx)?.into_array(); + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[25.0]); + Ok(()) +} + +#[rstest] +#[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] +#[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(InnerProduct); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("InnerProduct serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn inner_product_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} + +fn inner_product_vector_rhs() -> ArrayRef { + vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") +} + +fn inner_product_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn inner_product_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/l2_norm.rs b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs new file mode 100644 index 00000000000..38ce74d8c65 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::l2_norm::L2Norm; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. +fn eval_l2_norm(input: ArrayRef) -> VortexResult> { + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = L2Norm::new(); + let array = L2Norm::try_new_array(tensor_array(&[1], &[3.0])?)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + assert_close( + &eval_l2_norm(zero_width_vector_array::(3)?)?, + &[0.0, 0.0, 0.0], + ); + assert!(eval_l2_norm(vector_array(2, &[] as &[f64])?)?.is_empty()); + + let constant = Vector::constant_array::(&[], 3)?; + assert_close(&eval_l2_norm(constant)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[rstest] +#[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] +#[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] +#[case::single_element(&[1], &[7.0], &[7.0])] +#[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] +fn known_norms( + #[case] shape: &[usize], + #[case] elements: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let arr = tensor_array(shape, elements)?; + assert_close(&eval_l2_norm(arr)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let arr = tensor_array( + &[3], + &[ + 3.0, 4.0, 0.0, // norm = 5.0 + 0.0, 0.0, 0.0, // norm = 0.0 + 1.0, 1.0, 1.0, // norm = sqrt(3) + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); + Ok(()) +} + +#[test] +fn vector_multiple_rows() -> VortexResult<()> { + let arr = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // norm = 1.0 + 3.0, 4.0, 0.0, // norm = 5.0 + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 is masked as null. + let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: norm = 5.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is a non-null tensor should short-circuit to a +/// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so +/// execution stops at the [`Constant`] encoding instead of canonicalizing into a +/// [`PrimitiveArray`]. +#[test] +fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { + let input = literal_vector_array(&[3.0f64, 4.0], 4); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over a constant input must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// An extension array over constant storage is folded just like a top-level constant instead of +/// recomputing the same norm once per row. +#[test] +fn extension_backed_constant_yields_constant_output() -> VortexResult<()> { + let input = Vector::constant_array(&[3.0f64, 4.0], 4)?; + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over constant-backed extension storage must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of +/// the correct primitive dtype and length. +#[test] +fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { + let storage_dtype = DType::FixedSizeList( + DType::Primitive(PType::F64, Nullability::NonNullable).into(), + 2, + Nullability::Nullable, + ); + let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); + let null_scalar = Scalar::null(DType::Extension(ext_dtype)); + let input = ConstantArray::new(null_scalar, 3).into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("null constant input must produce a constant output"); + assert_eq!(constant.len(), 3); + assert!(constant.scalar().is_null()); + assert_eq!( + constant.dtype(), + &DType::Primitive(PType::F64, Nullability::Nullable) + ); + Ok(()) +} + +/// An `f32` column must dispatch at `f32` and produce an `f32` result, which is the property that +/// makes width polymorphism load-bearing rather than decorative. +#[rstest] +#[case::f32(&[3.0f32, 4.0], PType::F32)] +#[case::f64(&[3.0f64, 4.0], PType::F64)] +fn dispatches_at_input_width( + #[case] elements: &[T], + #[case] expected: PType, +) -> VortexResult<()> { + let arr = tensor_array(&[2], elements)?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L2Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + assert_eq!(prim.ptype(), expected); + Ok(()) +} + +/// `L2Norm(Normalized(normalized, norms))` reads back the authoritative stored norms rather than +/// recomputing over decoded coordinates. The normalized child here is deliberately *not* +/// unit-norm, mimicking lossy storage, so readthrough and recompute disagree: row 0 decodes to +/// `[6, 8]` (norm `10`) and row 1 to `[6, 0]` (norm `6`), while the stored norms are `5` and `2`. +#[test] +fn normalized_readthrough_returns_stored_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: A focused test of the lossy storage contract: the stored norms are authoritative + // even though this normalized child violates the unit-norm invariant. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) }.into_array(); + + assert_close(&eval_l2_norm(denorm)?, &[5.0, 2.0]); + Ok(()) +} + +/// The readthrough must survive a partially-null column. +/// +/// This pins the dense policy the row contract derives. Filtering could hand `reduce_encoded` a +/// filtered input, which is no longer an `ExactScalarFn`, silently falling back to +/// decode-and-recompute. For a lossy child that changes the answer: row 0 below would come back as +/// `10` (recomputed from `[6, 8]`) instead of the authoritative stored `5`. +#[test] +fn normalized_readthrough_survives_null_rows() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: Intentionally lossy, as in `normalized_readthrough_returns_stored_norms`, so that + // a recompute fallback is observable. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::from_iter([true, false])) } + .into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// The readthrough must still propagate validity carried by the `Normalized` parent. +#[test] +fn normalized_readthrough_propagates_parent_validity() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let denorm = Normalized::try_new( + normalized, + norms, + Validity::from_iter([true, false]), + &mut ctx, + )? + .into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +#[rstest] +#[case::fixed_shape_tensor(l2_norm_tensor_child())] +#[case::vector(l2_norm_vector_child())] +fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { + let original = L2Norm.try_new_array(child.len(), EmptyOptions, [child.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(L2Norm); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("L2Norm serialize must produce metadata"); + + let children = vec![child]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn l2_norm_tensor_child() -> ArrayRef { + tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") +} + +fn l2_norm_vector_child() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs new file mode 100644 index 00000000000..bb3726e9329 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for the tensor scalar functions. + +mod cosine_similarity; +mod inner_product; +mod l2_norm; +mod row; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs new file mode 100644 index 00000000000..8a7fa23edd4 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::Float; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::tensor_array; + +/// The marginal cost of a new tensor scalar function is this entire definition. Everything else +/// (null propagation, constants, validity, f16/f32/f64 dispatch, dtype checks, and constructors) is +/// derived. +#[derive(Clone, Debug, Default)] +struct L1Norm; + +impl RowFn for L1Norm { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.l1_norm"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l1_norm_row(row)) } + }) + }) + } +} + +fn l1_norm_row(row: &[T]) -> T { + row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) +} + +#[test] +fn derived_fn_executes_with_nulls() -> VortexResult<()> { + let arr = tensor_array(&[2], &[3.0, -4.0, 1.0, 1.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L1Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[7.0]); + Ok(()) +} + +/// A kernel written once serves every float width. +#[test] +fn derived_fn_dispatches_at_input_width() -> VortexResult<()> { + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + + let f32_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f32, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f32_result.ptype(), PType::F32); + assert_eq!(f32_result.as_slice::(), &[7.0f32]); + + let f64_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f64, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f64_result.ptype(), PType::F64); + Ok(()) +} diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 3e33fe20db9..73d98730092 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Shared helpers for tensor scalar functions. + use half::f16; +use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -22,6 +25,7 @@ use vortex_array::dtype::PType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -60,6 +64,19 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root } +/// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. +/// +/// L2 norm and cosine similarity share this implementation so prepared constant norms use the +/// same accumulation order as rows from per-row inputs. +pub(crate) fn l2_norm_row(row: &[T]) -> T { + let mut sum_squared = T::zero(); + for &element in row { + sum_squared = sum_squared + element * element; + } + + sum_squared.sqrt() +} + /// Extracts the `(normalized, norms)` children of a [`Normalized`]-encoded array. /// /// # Panics @@ -107,17 +124,20 @@ pub fn validate_tensor_float_input(input_dtype: &DType) -> VortexResult( - lhs: &'a DType, - rhs: &DType, -) -> VortexResult> { - vortex_ensure!( - lhs.eq_ignore_nullability(rhs), - "binary tensor expression expects inputs to have the same dtype, got {lhs} and {rhs}" - ); - validate_tensor_float_input(lhs) +/// Validates that every argument has the same float tensor dtype, ignoring nullability. +pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { + let (first, rest) = args + .split_first() + .ok_or_else(|| vortex_err!("tensor expression expects at least one input"))?; + + for arg in rest { + vortex_ensure!( + first.eq_ignore_nullability(arg), + "tensor expression expects inputs to have the same dtype, got {first} and {arg}" + ); + } + + validate_tensor_float_input(first) } /// The flat primitive elements of a tensor storage array, with typed row access. @@ -148,6 +168,23 @@ impl FlatElements { let slice = self.elems.as_slice::(); &slice[row_idx * self.list_size..][..self.list_size] } + + /// Returns the number of elements in each row. + #[must_use] + pub fn list_size(&self) -> usize { + self.list_size + } + + /// Returns the physical distance between rows, or zero when every row uses one stored value. + #[must_use] + pub fn row_stride(&self) -> usize { + if self.is_constant { 0 } else { self.list_size } + } + + /// Returns the elements as a typed buffer, performing the ptype check once for the batch. + pub fn into_buffer(self) -> Buffer { + self.elems.into_buffer::() + } } /// Extracts the flat primitive elements from a tensor storage array (FixedSizeList). @@ -285,7 +322,7 @@ impl BinaryTensorOpMetadata { let lhs_dtype = DType::from_proto(lhs_pb, session)?; let rhs_dtype = DType::from_proto(rhs_pb, session)?; - validate_binary_tensor_float_inputs(&lhs_dtype, &rhs_dtype)?; + validate_tensor_float_inputs(&[lhs_dtype.clone(), rhs_dtype.clone()])?; let lhs = children.get(0, &lhs_dtype, len)?; let rhs = children.get(1, &rhs_dtype, len)?; @@ -353,6 +390,18 @@ pub mod test_helpers { Vector::try_new_vector_array(flat_fsl(elements, dim)) } + /// Builds `rows` zero-width vectors over an empty typed element buffer. + pub fn zero_width_vector_array(rows: usize) -> VortexResult { + let storage = FixedSizeListArray::new( + Buffer::::empty().into_array(), + 0, + Validity::NonNullable, + rows, + ) + .into_array(); + Vector::try_new_vector_array(storage) + } + /// Builds a [`FixedShapeTensor`] extension array whose storage is a [`ConstantArray`], /// representing a single query tensor broadcast to `len` rows. pub fn constant_tensor_array>( diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index ad3b96d1bff..492bc837b89 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -35,11 +35,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -79,7 +81,7 @@ pub fn build_similarity_search_tree>( let num_rows = data.len(); let query_vec = Vector::constant_array(query, num_rows)?; - let cosine = CosineSimilarity::try_new_array(data, query_vec)?.into_array(); + let cosine = CosineSimilarity.try_new_array(num_rows, EmptyOptions, [data, query_vec])?; let threshold_scalar = Scalar::primitive(threshold, Nullability::NonNullable); let threshold_array = ConstantArray::new(threshold_scalar, num_rows).into_array();