test: read Lance tables from DuckDB (#10866)

* test: read Lance tables from DuckDB

The LanceDB and Spark suites go through the catalog. DuckDB does not: its
lance extension reaches the data over S3 with no namespace involved, which
exercises the other half of the design - a table bucket's layout is a
valid Lance dataset directory, so a table stays readable when the catalog
is not in the path.

    scan_rows=128
    scan_columns=id,title,vector
    filtered_rows=5
    nearest=1,0,2

It also pins the one place the layout costs us. DuckDB's replacement scan
recognises a dataset by a .lance path suffix, and tables created through
this catalog deliberately have none: the catalog entry is the dataset
directory, a table name may not contain a dot, and a suffix would leak
into ARNs and policies. So __lance_scan is the way in, and the bare
SELECT ... FROM 's3://...' form does not see these tables.

The test asserts both halves - a suffixed path is read, a suffix-less one
is not - so if the extension ever recognises a bare directory, it fails
and says to update the documentation rather than leaving it wrong.

Claude-Session: https://claude.ai/code/session_01Rkp1Mw5E89Jp6dzJFYiMrm

* test: require the catalog error from the suffix-less read

Any failure satisfied the old check - a missing extension, bad credentials,
an unreachable endpoint - so the assertion could pass without the
replacement scan ever classifying the path.

Claude-Session: https://claude.ai/code/session_01Rkp1Mw5E89Jp6dzJFYiMrm

* test: verify the Lance table bucket was actually created

weed shell prints a command's own failure and still exits 0, so the harness
would go on to blame DuckDB for a bucket that was never made.

Claude-Session: https://claude.ai/code/session_01Rkp1Mw5E89Jp6dzJFYiMrm

* test: bound the Docker probe

An unhealthy daemon makes docker version hang, and the probe runs before the
test has a timeout of its own.

Claude-Session: https://claude.ai/code/session_01Rkp1Mw5E89Jp6dzJFYiMrm

* test: order the aggregates the assertions read

string_agg over an unordered relation may return the names, and the vector
search's ids, in any order, so the expectations could fail on a run where
nothing changed.

Claude-Session: https://claude.ai/code/session_01Rkp1Mw5E89Jp6dzJFYiMrm

* test: do not persist credentials in the DuckDB Lance checkout

The job only uploads a log on failure; nothing in it pushes.

Claude-Session: https://claude.ai/code/session_01Rkp1Mw5E89Jp6dzJFYiMrm
This commit is contained in:
Chris Lu
2026-08-21 15:22:38 -07:00
committed by GitHub
parent 35d53a20f6
commit 5e7ab43ddd
6 changed files with 673 additions and 0 deletions
+75
View File
@@ -1013,6 +1013,81 @@ jobs:
path: test/s3tables/lifecycle/test-output.log
retention-days: 3
duckdb-lance-tests:
name: DuckDB Lance Integration Tests
runs-on: ubuntu-22.04
timeout-minutes: 30
steps:
- name: Check out code
uses: actions/checkout@v7
with:
# The job uploads a test log on failure; nothing here needs to push,
# so do not leave a token in the checkout for it to pick up.
persist-credentials: false
- name: Set up Go
uses: actions/setup-go@v7
with:
go-version-file: 'go.mod'
id: go
- name: Configure Docker Hub mirror
run: |
echo '{"registry-mirrors": ["https://mirror.gcr.io"]}' | sudo tee /etc/docker/daemon.json
sudo systemctl restart docker
- name: Pre-pull images
run: |
pull() { for i in 1 2 3; do docker pull "$1" && return 0; sleep 15; done; return 1; }
pull duckdb/duckdb:latest
pull python:3.11-slim
- name: Run go mod tidy
run: go mod tidy
- name: Build SeaweedFS
run: |
cd weed && go build -buildvcs=false .
- name: Run DuckDB Lance Integration Tests
timeout-minutes: 25
working-directory: test/s3tables/catalog_duckdb_lance
run: |
set -x
set -o pipefail
echo "=== System Information ==="
uname -a
free -h
df -h
docker info
echo "=== Starting DuckDB Lance Tests ==="
go test -v -timeout 20m . 2>&1 | tee test-output.log || {
echo "DuckDB Lance integration tests failed"
exit 1
}
- name: Show test output on failure
if: failure()
working-directory: test/s3tables/catalog_duckdb_lance
run: |
echo "=== Test Output ==="
if [ -f test-output.log ]; then
tail -200 test-output.log
fi
echo "=== Process information ==="
ps aux | grep -E "(weed|test|docker|duckdb)" || true
- name: Upload test logs on failure
if: failure()
uses: actions/upload-artifact@v7
with:
name: duckdb-lance-test-logs
path: test/s3tables/catalog_duckdb_lance/test-output.log
retention-days: 3
s3-tables-build-verification:
name: S3 Tables Build Verification
runs-on: ubuntu-22.04
@@ -0,0 +1,15 @@
# Seeds a Lance dataset so DuckDB has something to read. The catalog records
# where a table lives and does not carry its data, so the writing half is
# pylance; DuckDB then reads it over S3 with no catalog involved.
FROM python:3.11-slim
WORKDIR /app
RUN pip install --no-cache-dir \
"lance-namespace==0.8.6" \
"pylance==10.0.0" \
"pyarrow==25.0.1"
COPY seed_table.py /app/
CMD ["python3", "/app/seed_table.py", "--help"]
@@ -0,0 +1,66 @@
# DuckDB Lance Integration Test
Reads SeaweedFS Lance tables from DuckDB's `lance` core extension, the
counterpart of the DuckDB Iceberg tests in `../catalog/`.
## What this one proves that the others do not
The LanceDB and Spark suites go through the catalog. DuckDB does not: it reaches
the data over S3 with no namespace involved. That exercises the other half of the
design — a table bucket's layout is a valid Lance dataset directory, so a table
stays readable when the catalog is not in the path.
## What it does
`TestDuckDBLance`:
1. Starts a `weed mini` cluster with S3 and the Lance Namespace enabled.
2. Creates a table bucket declared `LANCE`, then declares a table through the
namespace and writes a dataset into it with pylance.
3. Runs `duckdb_lance_ops.sql` in `duckdb/duckdb:latest`.
| Step | What it proves |
| --- | --- |
| `__lance_scan(s3://…)` | DuckDB reads a table this catalog created |
| `DESCRIBE` | the schema survived, vector column included |
| `WHERE id < 5` | the filter path |
| `lance_vector_search` | vector search over data behind SeaweedFS |
| a `.lance` path | the replacement scan works on a suffixed path |
| a suffix-less path | and does **not** see one without the suffix |
## The `.lance` suffix
DuckDB's replacement scan — `SELECT * FROM 's3://…'` — recognises a Lance dataset
by a `.lance` path suffix. Tables created through this catalog deliberately have
none: the catalog entry *is* the dataset directory, table names may not contain a
dot, and a suffix would leak into ARNs and policies.
So from DuckDB, a table in a SeaweedFS table bucket is read with
`__lance_scan('s3://bucket/namespace/table')` rather than the bare `FROM 's3://…'`
form. The test asserts both halves, so if the extension ever recognises a
suffix-less directory the test fails and tells us to update the documentation.
## Credentials
```sql
CREATE SECRET seaweedfs (
TYPE lance,
ACCESS_KEY_ID '', SECRET_ACCESS_KEY '',
REGION 'us-east-1',
ENDPOINT 'http://seaweed:8333',
ALLOW_HTTP true,
VIRTUAL_HOSTED_STYLE_REQUEST false
);
```
Those are object_store's key names, the same as everywhere else Lance touches
storage.
## Running it
cd test/s3tables/catalog_duckdb_lance
(cd ../../../weed && go build .) # the harness runs this binary
go test -run TestDuckDBLance -v -timeout 30m .
Skipped without Docker, in `-short` mode, and if the DuckDB image cannot load the
extension.
@@ -0,0 +1,48 @@
-- Read a SeaweedFS Lance table from DuckDB.
--
-- The lance extension reaches the data over S3 rather than through the
-- namespace, so what is exercised here is the layout on the S3 door and the
-- credentials, not the catalog protocol. That is the point: this is the client
-- that proves a Lance table stays readable without a catalog at all.
--
-- Placeholders are substituted by the Go harness: __ENDPOINT__, __KEY__,
-- __SECRET__, __TABLE__ (s3://bucket/ns/table), __SUFFIXED__ (the same data at a
-- path ending in .lance).
INSTALL lance;
LOAD lance;
CREATE SECRET seaweedfs (
TYPE lance,
ACCESS_KEY_ID '__KEY__',
SECRET_ACCESS_KEY '__SECRET__',
REGION 'us-east-1',
ENDPOINT '__ENDPOINT__',
ALLOW_HTTP true,
VIRTUAL_HOSTED_STYLE_REQUEST false
);
-- 1. The table the catalog created, read by URI. A table bucket's layout is a
-- valid Lance dataset directory, which is what makes this possible.
SELECT 'scan_rows=' || count(*) FROM __lance_scan('__TABLE__');
-- 2. Its schema survived, vector column included. Ordered, because string_agg
-- over an unordered relation is free to return the names in any order.
SELECT 'scan_columns=' || string_agg(column_name, ',' ORDER BY column_name)
FROM (DESCRIBE SELECT * FROM __lance_scan('__TABLE__'));
-- 3. A filter, so it is not only a full scan.
SELECT 'filtered_rows=' || count(*) FROM __lance_scan('__TABLE__') WHERE id < 5;
-- 4. Vector search, which is what the format is for. No index is built here, so
-- this is a brute-force search; the ids nearest the query are what matters.
SELECT 'nearest=' || string_agg(id::VARCHAR, ',' ORDER BY _distance ASC, id ASC)
FROM lance_vector_search('__TABLE__', 'vector',
[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], k := 3);
-- 5. DuckDB's replacement scan recognises a path by its .lance suffix. Tables
-- this catalog creates deliberately have no suffix - the name is the table's
-- name, and a suffix would leak into ARNs and policies - so the bare
-- SELECT ... FROM 's3://…' form does not see them, and __lance_scan is the
-- way in. This asserts both halves so that a change upstream is noticed.
SELECT 'suffixed_rows=' || count(*) FROM '__SUFFIXED__';
@@ -0,0 +1,390 @@
// Package duckdblance reads SeaweedFS Lance tables from DuckDB's lance
// extension, the counterpart of the DuckDB Iceberg tests next door.
//
// DuckDB reaches the data over S3 rather than through the namespace, so what
// this proves is the other half of the design: a table bucket's layout is a
// valid Lance dataset directory, and a table stays readable with no catalog in
// the path at all. It also pins the one place that costs us - DuckDB's
// replacement scan recognises a dataset by its .lance suffix, which tables
// created through this catalog deliberately do not have.
package duckdblance
import (
"context"
"fmt"
"math/rand"
"net/http"
"os"
"os/exec"
"path/filepath"
"strings"
"testing"
"time"
"github.com/seaweedfs/seaweedfs/test/testutil"
)
const (
// Tracks whatever DuckDB users are actually running; the lance extension is
// a core extension, so there is nothing to pin beyond the image.
duckDBImage = "duckdb/duckdb:latest"
// Seeds a dataset for DuckDB to read.
seedImage = "seaweedfs-lance-seed"
startupTimeout = 60 * time.Second
clientTimeout = 15 * time.Minute
seededRows = 128
)
// TestDuckDBLance reads a table this catalog created, from DuckDB.
func TestDuckDBLance(t *testing.T) {
if testing.Short() {
t.Skip("Skipping integration test in short mode")
}
if !hasDocker() {
t.Skip("Docker not available, skipping DuckDB Lance integration test")
}
env := newEnvironment(t)
defer env.cleanup()
env.start(t)
bucket := "duckdb-" + randomSuffix()
env.createTableBucket(t, bucket)
requireDuckDBLance(t)
buildSeedImage(t)
env.seedTable(t, bucket)
env.runDuckDB(t, bucket)
}
type environment struct {
weedBinary string
dataDir string
bindIP string
masterPort int
masterGrpcPort int
volumePort int
volumeGrpcPort int
filerPort int
filerGrpcPort int
s3Port int
s3GrpcPort int
lancePort int
accessKey string
secretKey string
weedCancel context.CancelFunc
weedCmd *exec.Cmd
}
func newEnvironment(t *testing.T) *environment {
t.Helper()
wd, err := os.Getwd()
if err != nil {
t.Fatalf("get working directory: %v", err)
}
seaweedDir := wd
for i := 0; i < 5; i++ {
if _, err := os.Stat(filepath.Join(seaweedDir, "go.mod")); err == nil {
break
}
seaweedDir = filepath.Dir(seaweedDir)
}
weedBinary := filepath.Join(seaweedDir, "weed", "weed")
if info, statErr := os.Stat(weedBinary); statErr == nil && !info.IsDir() {
// `make test` builds first; a plain `go test` will otherwise drive a
// binary from days ago and report a pass for code it never ran.
t.Logf("using %s, built %s", weedBinary, info.ModTime().Format(time.RFC3339))
} else {
weedBinary = "weed"
if _, err := exec.LookPath(weedBinary); err != nil {
t.Skip("weed binary not found, skipping integration test")
}
}
dataDir, err := os.MkdirTemp("", "seaweed-duckdb-lance-test-*")
if err != nil {
t.Fatalf("create temp dir: %v", err)
}
ports := testutil.MustAllocatePorts(t, 9)
return &environment{
weedBinary: weedBinary,
dataDir: dataDir,
bindIP: testutil.FindBindIP(),
masterPort: ports[0],
masterGrpcPort: ports[1],
volumePort: ports[2],
volumeGrpcPort: ports[3],
filerPort: ports[4],
filerGrpcPort: ports[5],
s3Port: ports[6],
s3GrpcPort: ports[7],
lancePort: ports[8],
accessKey: "AKIAIOSFODNN7EXAMPLE",
secretKey: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
}
}
func (env *environment) start(t *testing.T) {
t.Helper()
iamConfigPath, err := testutil.WriteIAMConfig(env.dataDir, env.accessKey, env.secretKey)
if err != nil {
t.Fatalf("write IAM config: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
env.weedCancel = cancel
cmd := exec.CommandContext(ctx, env.weedBinary, "mini",
"-master.port", fmt.Sprintf("%d", env.masterPort),
"-master.port.grpc", fmt.Sprintf("%d", env.masterGrpcPort),
"-volume.port", fmt.Sprintf("%d", env.volumePort),
"-volume.port.grpc", fmt.Sprintf("%d", env.volumeGrpcPort),
"-filer.port", fmt.Sprintf("%d", env.filerPort),
"-filer.port.grpc", fmt.Sprintf("%d", env.filerGrpcPort),
"-s3.port", fmt.Sprintf("%d", env.s3Port),
"-s3.port.grpc", fmt.Sprintf("%d", env.s3GrpcPort),
"-s3.port.lance", fmt.Sprintf("%d", env.lancePort),
"-s3.config", iamConfigPath,
"-ip", env.bindIP,
"-ip.bind", "0.0.0.0",
"-dir", env.dataDir,
)
cmd.Dir = env.dataDir
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
cmd.Env = append(os.Environ(),
"AWS_ACCESS_KEY_ID="+env.accessKey,
"AWS_SECRET_ACCESS_KEY="+env.secretKey,
)
if err := cmd.Start(); err != nil {
t.Fatalf("start SeaweedFS: %v", err)
}
env.weedCmd = cmd
// The namespace answers /v1/table once it is serving, which is a cheaper
// readiness check than waiting on a bucket that does not exist yet.
url := fmt.Sprintf("http://%s:%d/v1/table", env.bindIP, env.lancePort)
if !waitForHTTP(url, startupTimeout) {
t.Fatalf("the Lance namespace did not become ready at %s", url)
}
}
// waitForHTTP polls until the URL answers at all. An auth refusal counts: it
// means the server is up, which is the only thing being waited on.
func waitForHTTP(url string, timeout time.Duration) bool {
client := &http.Client{Timeout: 2 * time.Second}
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
resp, err := client.Get(url)
if err != nil {
time.Sleep(500 * time.Millisecond)
continue
}
status := resp.StatusCode
resp.Body.Close()
if status < 500 {
return true
}
time.Sleep(500 * time.Millisecond)
}
return false
}
func (env *environment) cleanup() {
if env.weedCancel != nil {
env.weedCancel()
}
if env.weedCmd != nil {
_ = env.weedCmd.Wait()
}
if env.dataDir != "" {
_ = os.RemoveAll(env.dataDir)
}
}
// createTableBucket makes the bucket LanceDB will read through, declared LANCE
// so the catalog refuses anything of another format in it.
func (env *environment) createTableBucket(t *testing.T, bucket string) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
cmd := exec.CommandContext(ctx, env.weedBinary, "shell",
fmt.Sprintf("-master=%s:%d.%d", env.bindIP, env.masterPort, env.masterGrpcPort),
)
cmd.Stdin = strings.NewReader(fmt.Sprintf(
"s3tables.bucket -create -name %s -format LANCE -account 000000000000\nexit\n", bucket))
out, err := cmd.CombinedOutput()
if err != nil {
t.Fatalf("create table bucket %s: %v\n%s", bucket, err, out)
}
// weed shell reports a command's own failure on stdout and still exits 0, so
// the exit code alone would let a missing bucket through and turn a setup
// failure into a confusing engine failure later.
if !env.tableBucketExists(t, bucket) {
t.Fatalf("table bucket %s was not created:\n%s", bucket, out)
}
t.Logf("created LANCE table bucket %s", bucket)
}
// tableBucketExists asks the namespace, which lists table buckets at its root.
func (env *environment) tableBucketExists(t *testing.T, bucket string) bool {
t.Helper()
url := fmt.Sprintf("http://%s:%d/v1/namespace/%s/exists", env.bindIP, env.lancePort, bucket)
resp, err := http.Post(url, "application/json", strings.NewReader("{}"))
if err != nil {
t.Fatalf("ask the namespace whether %s exists: %v", bucket, err)
}
defer resp.Body.Close()
return resp.StatusCode == http.StatusOK
}
// runSpark runs the SQL driver inside the stock Spark image. The connector is
// pulled from Maven at submit time, the way the Iceberg Spark suite pulls its
// runtime, so nothing has to be built here.
// requireDuckDBLance skips rather than fails when the image cannot load the
// extension, so a DuckDB build without it reads as "not applicable" rather than
// as a broken catalog.
func requireDuckDBLance(t *testing.T) {
t.Helper()
const ready = "lance extension ready"
cmd := exec.Command("docker", "run", "--rm", "--entrypoint", "duckdb", duckDBImage,
"-c", fmt.Sprintf("INSTALL lance; LOAD lance; SELECT '%s' AS marker;", ready))
out, err := cmd.CombinedOutput()
if err != nil || !strings.Contains(string(out), ready) {
t.Skipf("DuckDB image cannot load the lance extension: %v\n%s", err, out)
}
}
func buildSeedImage(t *testing.T) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
defer cancel()
cmd := exec.CommandContext(ctx, "docker", "build", "-t", seedImage, "-f", "Dockerfile.seed", ".")
if out, err := cmd.CombinedOutput(); err != nil {
t.Fatalf("build the seed image: %v\n%s", err, out)
}
}
// seedTable declares a table through the namespace and writes a dataset into
// it, which is the split the catalog serves: it records where a table lives and
// does not carry its data.
func (env *environment) seedTable(t *testing.T, bucket string) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), clientTimeout)
defer cancel()
cmd := exec.CommandContext(ctx, "docker", "run", "--rm",
"--add-host", "host.docker.internal:host-gateway",
"-e", "AWS_ACCESS_KEY_ID="+env.accessKey,
"-e", "AWS_SECRET_ACCESS_KEY="+env.secretKey,
"-e", "AWS_REGION=us-east-1",
seedImage,
"python3", "/app/seed_table.py",
"--namespace-url", fmt.Sprintf("http://host.docker.internal:%d", env.lancePort),
"--s3-endpoint", fmt.Sprintf("http://host.docker.internal:%d", env.s3Port),
"--bucket", bucket,
"--rows", fmt.Sprintf("%d", seededRows),
"--access-key", env.accessKey,
"--secret-key", env.secretKey,
)
out, err := cmd.CombinedOutput()
t.Logf("seeder:\n%s", out)
if err != nil {
t.Fatalf("seeding the table failed: %v", err)
}
}
// runDuckDB substitutes the endpoint and paths into the SQL and runs it.
func (env *environment) runDuckDB(t *testing.T, bucket string) {
t.Helper()
sqlBytes, err := os.ReadFile("duckdb_lance_ops.sql")
if err != nil {
t.Fatalf("read the SQL: %v", err)
}
table := fmt.Sprintf("s3://%s/ml/embeddings", bucket)
sql := strings.NewReplacer(
"__ENDPOINT__", fmt.Sprintf("http://host.docker.internal:%d", env.s3Port),
"__KEY__", env.accessKey,
"__SECRET__", env.secretKey,
"__TABLE__", table,
"__SUFFIXED__", table+"-direct.lance",
).Replace(string(sqlBytes))
ctx, cancel := context.WithTimeout(context.Background(), clientTimeout)
defer cancel()
cmd := exec.CommandContext(ctx, "docker", "run", "--rm", "-i",
"--add-host", "host.docker.internal:host-gateway",
"--entrypoint", "duckdb", duckDBImage, "-c", sql)
out, err := cmd.CombinedOutput()
t.Logf("DuckDB output:\n%s", out)
if err != nil {
t.Fatalf("the DuckDB query failed: %v", err)
}
output := string(out)
for _, want := range []string{
fmt.Sprintf("scan_rows=%d", seededRows),
"scan_columns=id,title,vector",
"filtered_rows=5",
"nearest=1,0,2",
fmt.Sprintf("suffixed_rows=%d", seededRows),
} {
if !strings.Contains(output, want) {
t.Fatalf("DuckDB did not report %q", want)
}
}
// The other half of the suffix rule: a table this catalog created has no
// .lance suffix, so DuckDB's replacement scan does not see it and
// __lance_scan is the way in. If this ever starts working, the docs saying
// otherwise are wrong.
bare := exec.CommandContext(ctx, "docker", "run", "--rm", "-i",
"--add-host", "host.docker.internal:host-gateway",
"--entrypoint", "duckdb", duckDBImage,
"-c", fmt.Sprintf("INSTALL lance; LOAD lance; SELECT count(*) FROM '%s';", table))
bareOut, bareErr := bare.CombinedOutput()
// DuckDB exits nonzero for any error, so the exit status alone does not tell
// "the replacement scan refused the path" from "the query never ran": require
// the catalog error either way.
if !strings.Contains(string(bareOut), "does not exist") {
t.Fatalf("a suffix-less path did not fail the way the docs say it does (%v); "+
"if the replacement scan now reads it, update them:\n%s", bareErr, bareOut)
}
t.Logf("a suffix-less path is not seen by the replacement scan, as expected")
}
// hasDocker reports whether a Docker daemon answers. Bounded, because an
// unhealthy daemon makes `docker version` hang, and this runs before the test
// has a timeout of its own: better to skip than to eat the whole budget.
func hasDocker() bool {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
return exec.CommandContext(ctx, "docker", "version").Run() == nil
}
func randomSuffix() string {
const charset = "abcdefghijklmnopqrstuvwxyz0123456789"
suffix := make([]byte, 8)
for i := range suffix {
suffix[i] = charset[rand.Intn(len(charset))]
}
return string(suffix)
}
@@ -0,0 +1,79 @@
#!/usr/bin/env python3
"""Declare a table through the namespace and write a dataset into it.
Also writes the same rows to a second path whose name ends in `.lance`, because
DuckDB's replacement scan recognises a dataset by that suffix and tables created
through this catalog do not have one. The test asserts both behaviours so a
change upstream is noticed rather than silently making the docs wrong.
"""
import argparse
import sys
import warnings
warnings.filterwarnings("ignore")
import lance
import lance_namespace as ln
import pyarrow as pa
DIM = 8
def rows(count):
return pa.table(
{
"id": pa.array(list(range(count)), type=pa.int64()),
"title": pa.array([f"row-{i}" for i in range(count)]),
"vector": pa.array(
[[float(i) + d for d in range(DIM)] for i in range(count)],
type=pa.list_(pa.float32(), DIM),
),
}
)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--namespace-url", required=True)
parser.add_argument("--s3-endpoint", required=True)
parser.add_argument("--bucket", required=True)
parser.add_argument("--namespace", default="ml")
parser.add_argument("--table", default="embeddings")
parser.add_argument("--rows", type=int, default=128)
parser.add_argument("--access-key", default="any")
parser.add_argument("--secret-key", default="any")
args = parser.parse_args()
storage = {
"aws_endpoint": args.s3_endpoint,
"allow_http": "true",
"aws_access_key_id": args.access_key,
"aws_secret_access_key": args.secret_key,
"aws_region": "us-east-1",
}
ns = ln.connect("rest", {"uri": args.namespace_url})
ns.create_namespace(ln.CreateNamespaceRequest(id=[args.bucket], mode="EXIST_OK"))
ns.create_namespace(
ln.CreateNamespaceRequest(id=[args.bucket, args.namespace], mode="EXIST_OK")
)
declared = ns.declare_table(
ln.DeclareTableRequest(id=[args.bucket, args.namespace, args.table])
)
lance.write_dataset(rows(args.rows), declared.location, storage_options=storage,
mode="overwrite")
print(f"seeded {args.rows} rows at {declared.location}")
# The same data under a name DuckDB's replacement scan recognises. Written
# directly rather than declared, because a table name containing a dot is
# not a valid catalog name.
suffixed = f"s3://{args.bucket}/{args.namespace}/{args.table}-direct.lance"
lance.write_dataset(rows(args.rows), suffixed, storage_options=storage,
mode="overwrite")
print(f"seeded {args.rows} rows at {suffixed}")
return 0
if __name__ == "__main__":
sys.exit(main())