perf(ivfpq): speed up transposed PQ8 scans #304

Workflow file for this run

	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.

	name: CI

	on:
	push:
	branches:
	- main
	- release-*
	pull_request:
	branches:
	- main
	- release-*

	permissions:
	contents: read

	concurrency:
	group: ${{ github.workflow }}-${{ github.ref }}
	cancel-in-progress: true

	jobs:
	check:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v4

	- name: Set up minimum supported Python
	uses: actions/setup-python@v5
	with:
	python-version: '3.9'

	- name: Check benchmark tools on Python 3.9
	run: \|
	python3 -m py_compile \
	tools/convert_ann_benchmarks.py \
	tools/verify_binary_artifact.py

	- name: Validate .asf.yaml
	run: pip install pyyaml -q && python3 tools/validate_asf_yaml.py

	- name: Check ASF license headers
	run: python3 tools/check_license_headers.py

	- name: Check binary legal inventory
	run: python3 tools/verify_binary_artifact.py --source

	- name: Check reproducible source archive
	shell: bash
	run: \|
	set -euo pipefail
	first="$(mktemp)"
	second="$(mktemp)"
	tools/create_source_archive.sh 0.0.0 "$first"
	tools/create_source_archive.sh 0.0.0 "$second"
	cmp "$first" "$second"

	- name: Install cargo-deny
	uses: taiki-e/install-action@v2
	with:
	tool: cargo-deny@0.19.0

	- name: Check dependency licenses (Apache-compatible)
	run: cargo deny check licenses

	- name: Format
	run: cargo fmt --all -- --check

	- name: Clippy
	run: cargo clippy --all-targets --workspace -- -D warnings

	rust-test:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v4

	- name: Rust Cache
	uses: actions/cache@v4
	with:
	path: \|
	~/.cargo/registry
	~/.cargo/git
	target
	key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
	restore-keys: \|
	${{ runner.os }}-cargo-

	- name: Build
	run: cargo build --all-targets

	- name: Install cargo-llvm-cov
	uses: taiki-e/install-action@cargo-llvm-cov

	- name: Test with coverage
	run: cargo llvm-cov --workspace --summary-only

	c-test:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v4

	- name: Rust Cache
	uses: actions/cache@v4
	with:
	path: \|
	~/.cargo/registry
	~/.cargo/git
	target
	key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
	restore-keys: \|
	${{ runner.os }}-cargo-

	- name: Build FFI library
	run: cargo build --release -p paimon-vindex-ffi

	- name: Build C test
	run: \|
	cmake -S c -B c/build \
	-DPAIMON_VINDEX_FFI_LIB=${{ github.workspace }}/target/release/libpaimon_vindex_ffi.so
	cmake --build c/build

	- name: Run C test
	run: c/build/test_vindex
	env:
	LD_LIBRARY_PATH: ${{ github.workspace }}/target/release

	cpp-test:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v4

	- name: Rust Cache
	uses: actions/cache@v4
	with:
	path: \|
	~/.cargo/registry
	~/.cargo/git
	target
	key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
	restore-keys: \|
	${{ runner.os }}-cargo-

	- name: Build FFI library
	run: cargo build --release -p paimon-vindex-ffi

	- name: Build C++ test
	run: \|
	cmake -S cpp -B cpp/build \
	-DPAIMON_VINDEX_FFI_LIB=${{ github.workspace }}/target/release/libpaimon_vindex_ffi.so
	cmake --build cpp/build

	- name: Run C++ test
	run: cpp/build/test_vindex_cpp
	env:
	LD_LIBRARY_PATH: ${{ github.workspace }}/target/release

	jni-build:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v4

	- name: Rust Cache
	uses: actions/cache@v4
	with:
	path: \|
	~/.cargo/registry
	~/.cargo/git
	target
	key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
	restore-keys: \|
	${{ runner.os }}-cargo-

	- name: Set up JDK 8
	uses: actions/setup-java@v4
	with:
	java-version: '8'
	distribution: 'temurin'

	- name: Build JNI library
	run: cargo build -p paimon-vindex-jni --release

	- name: Test Java API
	run: >
	mvn -f java/pom.xml test
	-Dpaimon.vindex.native.path="${{ github.workspace }}/target/release/libpaimon_vindex_jni.so"

	- name: Test JNI native behavior
	run: \|
	java -cp java/target/test-classes:java/target/classes org.apache.paimon.index.vector.VectorIndexNativeValidationTest \
	"$(pwd)/target/release/libpaimon_vindex_jni.so"
	java -cp java/target/test-classes:java/target/classes org.apache.paimon.index.vector.VectorIndexNativePanicBoundaryTest \
	"$(pwd)/target/release/libpaimon_vindex_jni.so"
	java -cp java/target/test-classes:java/target/classes org.apache.paimon.index.vector.VectorIndexNativeHandleSafetyTest \
	"$(pwd)/target/release/libpaimon_vindex_jni.so"
	env:
	RAYON_NUM_THREADS: 4

	- name: Package and test the bundled JNI loader
	shell: bash
	run: \|
	set -euo pipefail
	mkdir -p java/src/main/resources/native/linux/x86_64
	cp target/release/libpaimon_vindex_jni.so \
	java/src/main/resources/native/linux/x86_64/
	mvn -f java/pom.xml package -DskipTests
	jar_file="java/target/paimon-vector-index-java-$(sed -n 's#.<version>$[^<]$</version>.*#\1#p' java/pom.xml \| sed -n '2p').jar"
	python3 tools/verify_binary_artifact.py \
	--jar "$jar_file" \
	--jar-native native/linux/x86_64/libpaimon_vindex_jni.so
	java -cp "$jar_file:java/target/test-classes" \
	org.apache.paimon.index.vector.VectorIndexNativeLoaderSmokeTest

	python-build:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v4

	- name: Rust Cache
	uses: actions/cache@v4
	with:
	path: \|
	~/.cargo/registry
	~/.cargo/git
	target
	key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
	restore-keys: \|
	${{ runner.os }}-cargo-

	- name: Set up Python
	uses: actions/setup-python@v5
	with:
	python-version: '3.12'

	- name: Build FFI library
	run: cargo build --release -p paimon-vindex-ffi

	- name: Test Python API
	working-directory: python
	run: \|
	pip install -e ".[test]"
	pytest -v
	env:
	PAIMON_VINDEX_LIB_PATH: ${{ github.workspace }}/target/release

	python-wheel:
	strategy:
	fail-fast: false
	matrix:
	os: [ubuntu-latest, macos-latest, windows-latest]
	runs-on: ${{ matrix.os }}
	steps:
	- uses: actions/checkout@v4

	- name: Rust Cache
	uses: actions/cache@v4
	with:
	path: \|
	~/.cargo/registry
	~/.cargo/git
	target
	key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
	restore-keys: \|
	${{ runner.os }}-cargo-

	- name: Set up Python
	uses: actions/setup-python@v5
	with:
	python-version: '3.12'

	- name: Build FFI library
	run: cargo build --release -p paimon-vindex-ffi

	- name: Build wheel
	working-directory: python
	run: \|
	python -m pip install --upgrade pip build
	python -m build --wheel

	- name: Verify wheel legal and native contents
	run: python tools/verify_binary_artifact.py --wheel "python/dist/*.whl"

	- name: Test wheel
	if: runner.os != 'Windows'
	working-directory: python
	run: \|
	python -m pip install dist/*.whl
	cd ..
	python - <<'PY'
	import io
	import numpy as np
	from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter

	class Input:
	def __init__(self, data):
	self.data = data

	def pread_many(self, ranges):
	return [self.data[pos : pos + length] for pos, length in ranges]

	data = np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [10.1, 10.0]], dtype=np.float32)
	ids = np.array([1, 2, 3, 4], dtype=np.int64)
	output = io.BytesIO()
	options = {"index.type": "ivf_flat", "dimension": "2", "nlist": "2", "metric": "l2"}
	training = VectorIndexTrainer.train(options, data)
	writer = VectorIndexWriter(training)
	writer.add_vectors(ids, data)
	writer.write(output)
	writer.close()

	reader = VectorIndexReader(Input(output.getvalue()))
	reader.optimize_for_search()
	result_ids, distances = reader.search(data[0], SearchParams.ivf(top_k=2, nprobe=2))
	assert result_ids[0] == 1, result_ids
	assert np.isfinite(distances[0]), distances
	reader.close()
	PY

	- name: Test wheel
	if: runner.os == 'Windows'
	working-directory: python
	shell: pwsh
	run: \|
	$wheel = Get-ChildItem dist/*.whl \| Select-Object -First 1
	python -m pip install $wheel.FullName
	Set-Location ..
	python -c @'
	import io
	import numpy as np
	from paimon_vindex import SearchParams, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter

	class Input:
	def __init__(self, data):
	self.data = data

	def pread_many(self, ranges):
	return [self.data[pos : pos + length] for pos, length in ranges]

	data = np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [10.1, 10.0]], dtype=np.float32)
	ids = np.array([1, 2, 3, 4], dtype=np.int64)
	output = io.BytesIO()
	options = {"index.type": "ivf_flat", "dimension": "2", "nlist": "2", "metric": "l2"}
	training = VectorIndexTrainer.train(options, data)
	writer = VectorIndexWriter(training)
	writer.add_vectors(ids, data)
	writer.write(output)
	writer.close()

	reader = VectorIndexReader(Input(output.getvalue()))
	reader.optimize_for_search()
	result_ids, distances = reader.search(data[0], SearchParams.ivf(top_k=2, nprobe=2))
	assert result_ids[0] == 1, result_ids
	assert np.isfinite(distances[0]), distances
	reader.close()
	'@

Navigation Menu

Uh oh!

perf(ivfpq): speed up transposed PQ8 scans #304

Workflow file

perf(ivfpq): speed up transposed PQ8 scans #304

Uh oh!

Workflow file for this run