diff --git a/docs/src/.pages b/docs/src/.pages new file mode 100644 index 0000000..ba6e6dc --- /dev/null +++ b/docs/src/.pages @@ -0,0 +1,6 @@ +nav: + - Welcome: index.md + - Install: install.md + - Config: config.md + - Performance: performance.md + - Operations: operations diff --git a/docs/src/config.md b/docs/src/config.md new file mode 100644 index 0000000..fc78309 --- /dev/null +++ b/docs/src/config.md @@ -0,0 +1,79 @@ +# Config + +Configuration options are grouped by area. The table connector uses `'connector' = 'lance'`; the +catalog types are `'lance'` (directory/S3) and `'lance-namespace'` (dir/rest). + +## Table connector options (`connector = 'lance'`) + +### Common + +| Option | Required | Default | Description | +|---|---|---|---| +| `path` | ✅ | — | Path to the Lance dataset | +| `hadoop.*` | ❌ | — | Prefix for Hadoop-family filesystem config (e.g. `hadoop.tbdsfs.meta`); stripped and injected into the Hadoop `Configuration` used for path resolution | + +### Read (Source) + +| Option | Required | Default | Description | +|---|---|---|---| +| `read.batch-size` | ❌ | 1024 | Read batch size | +| `read.limit` | ❌ | — | Maximum rows to read (limit pushdown) | +| `read.columns` | ❌ | — | Columns to read, comma separated | +| `read.filter` | ❌ | — | SQL `WHERE`-style filter predicate | +| `read.version` | ❌ | — | Time travel: read a specific dataset version | +| `read.as-of-timestamp` | ❌ | — | Time travel: read as of an ISO-8601 timestamp (ignored when `read.version` is set) | + +### Write (Sink) + +| Option | Required | Default | Description | +|---|---|---|---| +| `write.batch-size` | ❌ | 1024 | Write batch size | +| `write.mode` | ❌ | append | `append` or `overwrite` | +| `write.max-rows-per-file` | ❌ | 1000000 | Maximum rows per data file | + +### Vector index + +| Option | Required | Default | Description | +|---|---|---|---| +| `index.type` | ❌ | IVF_PQ | `IVF_PQ`, `IVF_HNSW`, or `IVF_FLAT` | +| `index.column` | ❌ | — | Vector column name to index | +| `index.num-partitions` | ❌ | 256 | IVF partition count | +| `index.num-sub-vectors` | ❌ | — | PQ sub-vector count (auto if unset) | +| `index.num-bits` | ❌ | 8 | PQ quantization bits (1–16) | +| `index.max-level` | ❌ | 7 | HNSW max level | +| `index.m` | ❌ | 16 | HNSW connections per level | +| `index.ef-construction` | ❌ | 100 | HNSW construction search width | + +### Vector search + +| Option | Required | Default | Description | +|---|---|---|---| +| `vector.column` | ❌ | — | Vector search column name | +| `vector.metric` | ❌ | L2 | `L2`, `Cosine`, or `Dot` | +| `vector.nprobes` | ❌ | 20 | IVF search probe count | +| `vector.ef` | ❌ | 100 | HNSW search width | +| `vector.refine-factor` | ❌ | — | Refine factor for recall | + +## Catalog options (`type = 'lance'`) + +Directory or S3 warehouse. + +| Option | Required | Default | Description | +|---|---|---|---| +| `warehouse` | ✅ | — | Warehouse path (local or `s3://…`) | +| `default-database` | ❌ | default | Default database name | +| `s3-access-key` | ❌ | — | S3 access key ID | +| `s3-secret-key` | ❌ | — | S3 secret access key | +| `s3-region` | ❌ | — | S3 region (e.g. `us-east-1`) | +| `s3-endpoint` | ❌ | — | S3 endpoint (for S3-compatible storage like MinIO) | +| `s3-virtual-hosted-style` | ❌ | true | Virtual-hosted-style URLs | +| `s3-allow-http` | ❌ | false | Allow HTTP (default HTTPS only) | + +## Namespace catalog options (`type = 'lance-namespace'`) + +| Option | Required | Default | Description | +|---|---|---|---| +| `impl` | ✅ | — | Namespace implementation: `dir` or `rest` | +| `root` | ❌ | — | Root path for directory namespace | +| `uri` | ❌ | — | URI for REST namespace | +| `default-database` | ❌ | default | Default database name | diff --git a/docs/src/index.md b/docs/src/index.md new file mode 100644 index 0000000..a52f225 --- /dev/null +++ b/docs/src/index.md @@ -0,0 +1,60 @@ +# Flink Lance Connector + +## Introduction + +The Apache Flink Connector for Lance allows Apache Flink to read and write datasets stored in the +[Lance](https://lance.org/) columnar format — an open lakehouse format optimized for multimodal AI +and vector search workloads. + +By using the Flink Connector for Lance, you can run Flink's stream/batch processing, SQL querying, +and stateful pipelines directly on Lance datasets, including native vector search. + +## Features + +The connector is built on the Flink Table API (`DynamicTableSource` / `DynamicTableSink`) plus +`CatalogFactory`. Specifically, you can use the Flink Connector for Lance to: + +* **Read & Write Lance Datasets**: append and overwrite datasets via Flink SQL or the DataStream API. +* **Column, Filter, Limit & Aggregate Pushdown**: push projections, `WHERE` predicates, limits and + aggregations down to Lance for efficient scans. +* **Vector Search**: KNN search over `ARRAY` columns with `L2`, `Cosine`, and `Dot` metrics, + via the `LanceVectorSearchFunction` table function. +* **Vector Index Building**: create `IVF_PQ`, `IVF_HNSW`, and `IVF_FLAT` indexes. +* **Time Travel**: read a historical version via `read.version` or `read.as-of-timestamp`. +* **Catalog Support**: a directory/S3 `lance` catalog and a `lance-namespace` catalog (dir / rest). + +## Quick Start + +Create a catalog and a table, then insert and query: + +```sql +-- Create a directory-based catalog +CREATE CATALOG lance_catalog WITH ( + 'type' = 'lance', + 'warehouse' = '/path/to/warehouse', + 'default-database' = 'default' +); + +USE CATALOG lance_catalog; + +-- Create a Lance table +CREATE TABLE vectors ( + id BIGINT, + content STRING, + embedding ARRAY +) WITH ( + 'connector' = 'lance', + 'path' = '/data/vectors', + 'write.batch-size' = '1024' +); + +-- Insert data +INSERT INTO vectors VALUES + (1, 'Hello World', ARRAY[0.1, 0.2, 0.3, 0.4]); + +-- Query data +SELECT * FROM vectors WHERE id > 0; +``` + +See [Install](install.md) for dependency setup and [Operations](operations/) for the full SQL +surface. diff --git a/docs/src/install.md b/docs/src/install.md new file mode 100644 index 0000000..c8b17a6 --- /dev/null +++ b/docs/src/install.md @@ -0,0 +1,49 @@ +# Install + +## Requirements + +* JDK 11 or higher +* Maven 3.6+ +* Apache Flink 1.18 / 1.19 / 1.20 + +The connector ships one artifact per supported Flink minor version: + +| Flink version | Artifact | +|---|---| +| 1.18 | `lance-flink-1.18` | +| 1.19 | `lance-flink-1.19` | +| 1.20 | `lance-flink-1.20` | + +## Dependencies + +The connector depends on `org.lance:lance-core` (7.0.0) and Apache Arrow (18.3.0). These are pulled +in transitively; you only need to add the connector artifact for your Flink version. + +Lance's Java bindings ship a platform-specific JNI native library +(`liblance_jni.so` / `liblance_jni.dylib`) inside the `lance-core` jar. Ensure you run on a +supported platform (linux-x86-64, linux-aarch64, darwin-aarch64). + +## Maven + +```xml + + org.apache.flink + lance-flink-1.18 + 0.1.0 + +``` + +## Build from source + +```bash +mvn clean verify +``` + +The build produces a fat jar per module (e.g. `lance-flink-1.18/target/lance-flink-1.18-*.jar`). +Add the jar to your Flink cluster or job classpath, then use the `lance` / `lance-namespace` +catalog types in SQL. + +## Note on Arrow and Netty + +The Arrow allocator defaults are set at runtime; if you see classloader-related SPI issues in +tests, set the system property `arrow.memory.allocator.type=Netty`. diff --git a/docs/src/operations/.pages b/docs/src/operations/.pages new file mode 100644 index 0000000..582cf82 --- /dev/null +++ b/docs/src/operations/.pages @@ -0,0 +1,4 @@ +nav: + - DDL: ddl + - DQL: dql + - DML: dml diff --git a/docs/src/operations/ddl/.pages b/docs/src/operations/ddl/.pages new file mode 100644 index 0000000..b95d1c3 --- /dev/null +++ b/docs/src/operations/ddl/.pages @@ -0,0 +1,3 @@ +nav: + - CREATE CATALOG: create-catalog.md + - CREATE TABLE: create-table.md diff --git a/docs/src/operations/ddl/create-catalog.md b/docs/src/operations/ddl/create-catalog.md new file mode 100644 index 0000000..ee2e5a0 --- /dev/null +++ b/docs/src/operations/ddl/create-catalog.md @@ -0,0 +1,80 @@ +# CREATE CATALOG + +The Lance Flink connector ships two catalog types, registered via SPI: + +| `type` | Class | Description | +|---|---|---| +| `lance` | `LanceCatalogFactory` | Directory-based catalog over a warehouse path (local or S3) | +| `lance-namespace` | `LanceNamespaceCatalogFactory` | Catalog backed by a Lance namespace (dir or REST) | + +## Directory catalog (`type = 'lance'`) + +```sql +CREATE CATALOG lance_catalog WITH ( + 'type' = 'lance', + 'warehouse' = '/path/to/warehouse', + 'default-database' = 'default' +); + +USE CATALOG lance_catalog; +``` + +### S3 warehouse + +```sql +CREATE CATALOG lance_s3_catalog WITH ( + 'type' = 'lance', + 'warehouse' = 's3://bucket-name/warehouse', + 'default-database' = 'default', + 's3-access-key' = 'your-access-key', + 's3-secret-key' = 'your-secret-key', + 's3-region' = 'us-east-1', + 's3-endpoint' = 'https://s3.amazonaws.com' +); +``` + +| Option | Required | Default | Description | +|---|---|---|---| +| `warehouse` | ✅ | — | Warehouse path (local or `s3://`) | +| `default-database` | ❌ | `default` | Default database | +| `s3-access-key` | ❌ | — | S3 access key | +| `s3-secret-key` | ❌ | — | S3 secret key | +| `s3-region` | ❌ | — | S3 region | +| `s3-endpoint` | ❌ | — | S3 endpoint (for MinIO etc.) | +| `s3-virtual-hosted-style` | ❌ | `true` | Virtual-hosted-style URLs | +| `s3-allow-http` | ❌ | `false` | Allow HTTP | + +## Namespace catalog (`type = 'lance-namespace'`) + +```sql +-- Directory-based namespace +CREATE CATALOG my_lance WITH ( + 'type' = 'lance-namespace', + 'impl' = 'dir', + 'root' = '/tmp/lance-warehouse' +); + +-- REST-based namespace +CREATE CATALOG my_lance WITH ( + 'type' = 'lance-namespace', + 'impl' = 'rest', + 'uri' = 'http://localhost:8080' +); +``` + +| Option | Required | Default | Description | +|---|---|---|---| +| `impl` | ✅ | — | `dir` or `rest` | +| `root` | ❌ | — | Root path for `dir` impl | +| `uri` | ❌ | — | URI for `rest` impl | +| `default-database` | ❌ | `default` | Default database | + +## Supported DDL + +| Statement | Status | +|---|---| +| `CREATE DATABASE` / `DROP DATABASE` / `ALTER DATABASE` | ✅ | +| `SHOW DATABASES` / `SHOW TABLES` | ✅ | +| `CREATE TABLE` / `DROP TABLE` / `RENAME TABLE` | ✅ | +| `ALTER TABLE` | ❌ — not supported (structure immutable) | +| `CREATE INDEX` | ❌ — no SQL DDL; configure `index.*` on the table | diff --git a/docs/src/operations/ddl/create-table.md b/docs/src/operations/ddl/create-table.md new file mode 100644 index 0000000..54702d5 --- /dev/null +++ b/docs/src/operations/ddl/create-table.md @@ -0,0 +1,56 @@ +# CREATE TABLE + +Lance tables are created through the dynamic table factory (`connector = 'lance'`). +The actual Lance dataset is created lazily on first write. + +## Minimal example + +```sql +CREATE TABLE vectors ( + id BIGINT, + content STRING, + embedding ARRAY +) WITH ( + 'connector' = 'lance', + 'path' = '/data/vectors' +); +``` + +## With a vector index + +```sql +CREATE TABLE doc_embeddings ( + doc_id BIGINT, + title STRING, + embedding ARRAY +) WITH ( + 'connector' = 'lance', + 'path' = '/data/embeddings', + 'index.type' = 'IVF_PQ', + 'index.column' = 'embedding', + 'index.num-partitions' = '256', + 'index.num-sub-vectors' = '16', + 'vector.metric' = 'COSINE' +); +``` + +## Required options + +| Option | Description | +|---|---| +| `path` | Path to the Lance dataset | + +## Type mapping + +| Lance / Arrow type | Flink type | +|---|---| +| Int8 / Int16 / Int32 / Int64 | TINYINT / SMALLINT / INT / BIGINT | +| Float32 / Float64 | FLOAT / DOUBLE | +| String | STRING | +| Boolean | BOOLEAN | +| Binary | BYTES | +| Date32 | DATE | +| Timestamp | TIMESTAMP | +| FixedSizeList\ | ARRAY\ | + +See [Config](../../config.md) for the full option reference. diff --git a/docs/src/operations/dml/.pages b/docs/src/operations/dml/.pages new file mode 100644 index 0000000..e9a5257 --- /dev/null +++ b/docs/src/operations/dml/.pages @@ -0,0 +1,2 @@ +nav: + - INSERT INTO: insert-into.md diff --git a/docs/src/operations/dml/insert-into.md b/docs/src/operations/dml/insert-into.md new file mode 100644 index 0000000..861a656 --- /dev/null +++ b/docs/src/operations/dml/insert-into.md @@ -0,0 +1,39 @@ +# INSERT INTO + +The Lance Flink sink appends rows to a Lance dataset. Write mode is controlled by +the `write.mode` option. + +## Write modes + +| `write.mode` | Behaviour | +|---|---| +| `append` (default) | Append rows to the existing dataset | +| `overwrite` | Replace the dataset on first write | + +## Example + +```sql +INSERT INTO vectors VALUES + (1, 'Hello World', ARRAY[0.1, 0.2, 0.3, 0.4]); +``` + +## Sink options + +| Option | Default | Description | +|---|---|---| +| `write.batch-size` | 1024 | Rows buffered before a flush | +| `write.mode` | `append` | `append` or `overwrite` | +| `write.max-rows-per-file` | 1000000 | Rows per data file | + +## Current limitations + +| Statement | Status | +|---|---| +| `INSERT INTO` (append) | ✅ | +| `INSERT OVERWRITE` | ✅ | +| `UPDATE` | ❌ — not implemented | +| `DELETE` | ❌ — in progress (see issue #63 / #74) | +| Primary key / upsert | ❌ — PK declaration and CDC changelog not yet supported | + +> The sink currently declares insert-only changelog mode. CDC `UPDATE` / `DELETE` +> support is tracked in the connector roadmap. diff --git a/docs/src/operations/dql/.pages b/docs/src/operations/dql/.pages new file mode 100644 index 0000000..35f1b23 --- /dev/null +++ b/docs/src/operations/dql/.pages @@ -0,0 +1,4 @@ +nav: + - SELECT: select.md + - Vector Search: vector-search.md + - Time Travel: time-travel.md diff --git a/docs/src/operations/dql/select.md b/docs/src/operations/dql/select.md new file mode 100644 index 0000000..efbc7ac --- /dev/null +++ b/docs/src/operations/dql/select.md @@ -0,0 +1,48 @@ +# SELECT + +The Lance Flink connector supports reading Lance datasets via Flink SQL `SELECT`. +Read optimizations are pushed down to Lance natively to reduce I/O. + +## Supported pushdowns + +| Ability | Interface | Notes | +|---|---|---| +| Column projection | `SupportsProjectionPushDown` | Only projected columns are read | +| Predicate (filter) | `SupportsFilterPushDown` | `WHERE` clauses are pushed to Lance | +| Limit | `SupportsLimitPushDown` | `LIMIT` is pushed down | +| Aggregate | `SupportsAggregatePushDown` | Eligible aggregates run natively | + +## Example + +```sql +-- Projection + filter + limit are all pushed down +SELECT id, content +FROM lance_table +WHERE id > 100 +LIMIT 10; +``` + +## Static read options + +The same read behaviour can be configured declaratively on the table DDL +without relying on planner pushdown: + +```sql +CREATE TABLE lance_table ( + id BIGINT, + content STRING, + embedding ARRAY +) WITH ( + 'connector' = 'lance', + 'path' = '/data/vectors', + 'read.columns' = 'id,content', + 'read.filter' = 'id > 100', + 'read.limit' = '10' +); +``` + +| Option | Description | +|---|---| +| `read.columns` | Comma-separated columns to read | +| `read.filter` | SQL `WHERE`-style filter string | +| `read.limit` | Maximum rows to read | diff --git a/docs/src/operations/dql/time-travel.md b/docs/src/operations/dql/time-travel.md new file mode 100644 index 0000000..5e9efea --- /dev/null +++ b/docs/src/operations/dql/time-travel.md @@ -0,0 +1,44 @@ +# Time Travel + +The Lance Flink connector supports reading historical versions of a dataset via +time-travel options, resolved uniformly through `LanceOpener`. + +## Options + +| Option | Type | Description | +|---|---|---| +| `read.version` | LONG | Read the given dataset version (highest precedence) | +| `read.as-of-timestamp` | STRING | Read as of an ISO-8601 timestamp; resolves to the newest version whose creation time is ≤ the timestamp | + +When both are set, `read.version` takes precedence. + +## Examples + +```sql +-- Read a specific version +SELECT * FROM lance_table /*+ OPTIONS('read.version' = '3') */; + +-- Read as of a timestamp (resolves to the newest version <= the timestamp) +SELECT * FROM lance_table /*+ OPTIONS('read.as-of-timestamp' = '2026-07-01T00:00:00Z') */; +``` + +Declaratively on the table DDL: + +```sql +CREATE TABLE lance_table ( + id BIGINT, + content STRING +) WITH ( + 'connector' = 'lance', + 'path' = '/data/vectors', + 'read.version' = '3' +); +``` + +## Semantics + +- `read.version` opens exactly that version. +- `read.as-of-timestamp` accepts any string parseable by + `Instant.parse`, `OffsetDateTime.parse`, or `ZonedDateTime.parse`; + a bare `yyyy-MM-ddTHH:mm:ss` is assumed UTC. +- A timestamp predating the oldest version raises a clear error. diff --git a/docs/src/operations/dql/vector-search.md b/docs/src/operations/dql/vector-search.md new file mode 100644 index 0000000..a0155b0 --- /dev/null +++ b/docs/src/operations/dql/vector-search.md @@ -0,0 +1,76 @@ +# Vector Search + +The Lance Flink connector exposes a table function for KNN vector search over a +Lance dataset, powered by Lance's IVF / HNSW indexes. + +## Function signature + +The function is `LanceVectorSearchFunction` — a Flink `TableFunction` registered as a +temporary function before use: + +```sql +CREATE TEMPORARY FUNCTION vector_search AS + 'org.apache.flink.connector.lance.table.LanceVectorSearchFunction' + LANGUAGE JAVA USING JAR '/path/to/lance-flink-1.18-0.1.0.jar'; +``` + +The registered function name (here `vector_search`) is user-chosen. Its `eval` overloads are: + +``` +vector_search(dataset_path, column_name, query_vector [, k [, metric]]) +``` + +- `dataset_path` — path to the Lance dataset +- `column_name` — the vector column to search +- `query_vector` — the query vector (`ARRAY`; also accepts `DECIMAL[]` / `DOUBLE[]` / `float[]`) +- `k` — number of nearest neighbours (default `10`) +- `metric` — distance metric (`L2` / `Cosine` / `Dot`, default `L2`) + +The emitted row is the source row plus a `_distance DOUBLE` column. + +## Distance metrics + +| Metric | Description | Range | +|---|---|---| +| `L2` | Euclidean distance | [0, ∞) | +| `Cosine` | Cosine similarity | [-1, 1] | +| `Dot` | Inner product | (-∞, ∞) | + +## Search options + +Configured on the table DDL or via `LanceOptions`: + +| Option | Default | Description | +|---|---|---| +| `vector.column` | — | Vector column name | +| `vector.metric` | `L2` | Distance metric | +| `vector.nprobes` | 20 | IVF probe count | +| `vector.ef` | 100 | HNSW search width | +| `vector.refine-factor` | — | Re-rank factor for recall | + +## Example + +```sql +CREATE TABLE vectors ( + id BIGINT, + content STRING, + embedding ARRAY +) WITH ( + 'connector' = 'lance', + 'path' = '/data/vectors', + 'index.type' = 'IVF_PQ', + 'index.column' = 'embedding', + 'vector.metric' = 'COSINE' +); + +-- KNN search for the 10 nearest vectors +SELECT * +FROM vectors, + LATERAL TABLE(vector_search( + '/data/vectors', + 'embedding', + ARRAY[0.1, 0.2, 0.3, 0.4], + 10, + 'COSINE' + )); +``` diff --git a/docs/src/performance.md b/docs/src/performance.md new file mode 100644 index 0000000..6bc839f --- /dev/null +++ b/docs/src/performance.md @@ -0,0 +1,49 @@ +# Performance + +This page covers vector index selection and the tuning knobs that most affect +read/search throughput and write amplification in the Lance Flink connector. + +## Vector index types + +| Index type | Best for | Memory | Recall | Key parameters | +|---|---|---|---|---| +| `IVF_FLAT` | Small datasets (< 100K vectors), exact-ish search | High | Highest | `index.num-partitions` | +| `IVF_PQ` | Large datasets, memory-constrained | Lowest | Good | `index.num-partitions`, `index.num-sub-vectors`, `index.num-bits` | +| `IVF_HNSW` | High recall, fast query latency | High | High | `index.num-partitions`, `index.m`, `index.ef-construction` | + +## Index selection guide + +| Scenario | Recommended index | Reason | +|---|---|---| +| < 100K vectors | `IVF_FLAT` | Highest accuracy, acceptable latency | +| 100K – 10M vectors | `IVF_PQ` | Good accuracy/memory trade-off | +| > 10M vectors | `IVF_PQ` (tuned) | Tune `index.num-partitions` and `index.num-sub-vectors` | +| High recall required | `IVF_HNSW` | Best accuracy, higher memory | +| Memory constrained | `IVF_PQ` | Most memory efficient | +| Real-time search | `IVF_HNSW` | Fastest query latency | + +## Search tuning + +| Option | Default | Effect | +|---|---|---| +| `vector.nprobes` | 20 | Number of IVF partitions probed per query. Higher = better recall, slower query | +| `vector.ef` | 100 | HNSW search width. Higher = better recall, slower query | +| `vector.refine-factor` | — | Refines top-k results by re-ranking candidates. Higher = better recall | + +## Write amplification + +| Option | Default | Effect | +|---|---|---| +| `write.batch-size` | 1024 | Rows buffered before a flush; larger = fewer, bigger writes | +| `write.max-rows-per-file` | 1000000 | Rows per data file; larger = fewer files, larger compaction units | + +## Distance metrics + +| Metric | Description | Range | +|---|---|---| +| `L2` | Euclidean distance | [0, ∞) | +| `Cosine` | Cosine similarity | [-1, 1] | +| `Dot` | Inner product | (-∞, ∞) | + +> **Note:** No benchmark numbers are published yet. These knobs are the exposed +> surface; measure against your own workload to tune them. diff --git a/src/main/java/org/apache/flink/connector/lance/ArrowArrayStreams.java b/src/main/java/org/apache/flink/connector/lance/ArrowArrayStreams.java new file mode 100644 index 0000000..eb0d671 --- /dev/null +++ b/src/main/java/org/apache/flink/connector/lance/ArrowArrayStreams.java @@ -0,0 +1,121 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.arrow.c.ArrowArrayStream; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.types.pojo.Schema; + +import org.lance.Dataset; +import org.lance.merge.MergeInsertParams; +import org.lance.merge.MergeInsertResult; + +import java.io.IOException; + +/** + * Bridges a single {@link VectorSchemaRoot} to the Arrow C Data Interface + * ({@link ArrowArrayStream}) required by {@link Dataset#mergeInsert}. + * + *

The Arrow Java SDK does not expose a public {@code VectorSchemaRoot -> ArrowArrayStream} + * adapter, so this class implements the minimal {@link ArrowReader} that yields exactly one + * record batch and exports it via {@link Data#exportArrayStream}. The exported stream lazily + * consumes the reader (whose lifecycle is owned by the stream's release callback); the record + * batch itself is closed by the caller here once the merge completes. + */ +public final class ArrowArrayStreams { + + private ArrowArrayStreams() { + // utility class + } + + /** + * Execute a native merge-insert (upsert) of a single batch. + * + * @param dataset target Lance dataset + * @param params merge-insert semantics (on keys, matched/unmatched behavior) + * @param allocator Arrow allocator (must stay open for the duration of the call) + * @param root populated data batch + * @return the merge-insert result + */ + public static MergeInsertResult mergeInsert( + Dataset dataset, + MergeInsertParams params, + BufferAllocator allocator, + VectorSchemaRoot root) throws IOException { + ArrowRecordBatch batch = new VectorUnloader(root).getRecordBatch(); + ArrowReader reader = new SingleBatchReader(allocator, root.getSchema(), batch); + ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); + try { + Data.exportArrayStream(allocator, reader, stream); + return dataset.mergeInsert(params, stream); + } finally { + // Closing the stream triggers its release callback, which closes the reader + // (and therefore the reader-owned VectorSchemaRoot). + stream.close(); + batch.close(); + } + } + + /** + * An {@link ArrowReader} that yields exactly one record batch and then reports end-of-stream. + */ + private static final class SingleBatchReader extends ArrowReader { + + private final Schema schema; + private final ArrowRecordBatch batch; + private boolean consumed; + + SingleBatchReader(BufferAllocator allocator, Schema schema, ArrowRecordBatch batch) { + super(allocator); + this.schema = schema; + this.batch = batch; + this.consumed = false; + } + + @Override + public boolean loadNextBatch() throws IOException { + if (consumed) { + return false; + } + consumed = true; + loadRecordBatch(batch); + return true; + } + + @Override + public long bytesRead() { + return 0L; + } + + @Override + protected void closeReadSource() throws IOException { + // nothing to release here; the record batch is owned by the enclosing caller + } + + @Override + protected Schema readSchema() throws IOException { + return schema; + } + } +} diff --git a/src/main/java/org/apache/flink/connector/lance/LanceSink.java b/src/main/java/org/apache/flink/connector/lance/LanceSink.java index 84e3069..ef3757e 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceSink.java @@ -156,7 +156,14 @@ public void flush() throws IOException { converter.toVectorSchemaRoot(buffer, root); String datasetPath = options.getPath(); - + + // A peer subtask may have created the dataset since this sink opened (multi-subtask + // first write). Re-check existence BEFORE Fragment.write below, which itself creates + // the dataset data directory and would otherwise make Files.exists unreliable. + if (!datasetExists && Files.exists(Paths.get(datasetPath))) { + datasetExists = true; + } + // Build write parameters WriteParams writeParams = new WriteParams.Builder() .withMaxRowsPerFile(options.getWriteMaxRowsPerFile()) diff --git a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java new file mode 100644 index 0000000..cb67c46 --- /dev/null +++ b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java @@ -0,0 +1,379 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.connector.lance.config.LanceOptions; +import org.apache.flink.connector.lance.converter.LanceTypeConverter; +import org.apache.flink.connector.lance.converter.RowDataConverter; +import org.apache.flink.runtime.state.FunctionInitializationContext; +import org.apache.flink.runtime.state.FunctionSnapshotContext; +import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction; +import org.apache.flink.streaming.api.functions.sink.RichSinkFunction; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.BooleanType; +import org.apache.flink.table.types.logical.DoubleType; +import org.apache.flink.table.types.logical.FloatType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.SmallIntType; +import org.apache.flink.table.types.logical.TinyIntType; +import org.apache.flink.table.types.logical.VarCharType; +import org.apache.flink.types.RowKind; + +import org.lance.CommitBuilder; +import org.lance.Dataset; +import org.lance.Fragment; +import org.lance.FragmentMetadata; +import org.lance.Transaction; +import org.lance.WriteParams; +import org.lance.merge.MergeInsertParams; +import org.lance.operation.Overwrite; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Keyed sink for Lance tables declared with a primary key. + * + *

Unlike {@link LanceSink} (append-only), this sink supports the CDC changelog kinds + * {@code +I}/{@code +U}/{-D}. It relies on {@code DataStream#keyBy} upstream to guarantee that all + * events for a given primary key arrive at the same subtask in order; it then collapses the + * buffered events per key to a single final action and applies it at checkpoint boundaries: + * + *

    + *
  • {@code INSERT}/{@code UPDATE_AFTER} → native upsert via + * {@link Dataset#mergeInsert} ({@code WhenMatched.UpdateAll} + + * {@code WhenNotMatched.InsertAll}).
    • + *
  • {@code DELETE} → {@link Dataset#delete} with an OR-of-AND predicate.
    • + *
  • {@code UPDATE_BEFORE} → dropped (upsert has no need for the old value).
    • + *
+ */ +public class LanceUpsertSink extends RichSinkFunction implements CheckpointedFunction { + + private static final long serialVersionUID = 1L; + private static final Logger LOG = LoggerFactory.getLogger(LanceUpsertSink.class); + + private final LanceOptions options; + private final RowType rowType; + private final List primaryKeys; + private final int[] keyIndices; + + private transient BufferAllocator allocator; + private transient Dataset dataset; + private transient RowDataConverter converter; + private transient Schema arrowSchema; + private transient Map buffer; + private transient long totalWrittenRows; + + public LanceUpsertSink(LanceOptions options, RowType rowType, List primaryKeys, int[] keyIndices) { + this.options = options; + this.rowType = rowType; + this.primaryKeys = primaryKeys; + this.keyIndices = keyIndices; + } + + @Override + public void open(Configuration parameters) throws Exception { + super.open(parameters); + + LOG.info("Opening Lance Upsert Sink: {}", options.getPath()); + + this.allocator = new RootAllocator(Long.MAX_VALUE); + this.buffer = new LinkedHashMap<>(); + this.totalWrittenRows = 0; + this.converter = new RowDataConverter(rowType); + this.arrowSchema = LanceTypeConverter.toArrowSchema(rowType); + + String datasetPath = options.getPath(); + if (datasetPath == null || datasetPath.isEmpty()) { + throw new IllegalArgumentException("Lance dataset path cannot be empty"); + } + + Path path = Paths.get(datasetPath); + boolean datasetExists = Files.exists(path); + + if (datasetExists && options.getWriteMode() == LanceOptions.WriteMode.OVERWRITE) { + LOG.info("Overwrite mode, deleting existing dataset: {}", datasetPath); + deleteDirectory(path); + datasetExists = false; + } + + if (datasetExists) { + this.dataset = Dataset.open(datasetPath, allocator); + } + + LOG.info("Lance Upsert Sink opened, primary keys: {}", primaryKeys); + } + + @Override + public void invoke(RowData value, Context context) { + RowKind kind = value.getRowKind(); + switch (kind) { + case INSERT: + case UPDATE_AFTER: + buffer.put(extractKey(value), value); + break; + case DELETE: + buffer.put(extractKey(value), value); + break; + case UPDATE_BEFORE: + // upsert has no use for the old value + break; + default: + LOG.warn("Ignoring unsupported RowKind: {}", kind); + } + } + + /** + * Flush the collapsed per-key buffer to Lance. + */ + public void flush() throws IOException { + if (buffer.isEmpty()) { + return; + } + + List upserts = new ArrayList<>(); + List deletes = new ArrayList<>(); + for (RowData row : buffer.values()) { + if (row.getRowKind() == RowKind.DELETE) { + deletes.add(row); + } else { + upserts.add(row); + } + } + + if (dataset == null) { + // First write: create the dataset from upserts only (deletes have no target yet). + if (upserts.isEmpty()) { + buffer.clear(); + return; + } + createDataset(upserts); + PrimaryKeyPersistence.persist(dataset, primaryKeys); + totalWrittenRows += upserts.size(); + } else { + if (!upserts.isEmpty()) { + mergeInsertRows(upserts); + } + if (!deletes.isEmpty()) { + deleteRows(deletes); + } + totalWrittenRows += upserts.size() + deletes.size(); + } + + buffer.clear(); + } + + /** + * Create the dataset on first write (equivalent to {@code INSERT} into an empty target). + */ + private void createDataset(List rows) throws IOException { + String datasetPath = options.getPath(); + + // A peer subtask may have created the dataset since this sink opened (multi-subtask first + // write). Falling back to merge-insert avoids clobbering its data with Overwrite. + if (Files.exists(Paths.get(datasetPath))) { + this.dataset = Dataset.open(datasetPath, allocator); + mergeInsertRows(rows); + return; + } + + try (VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema, allocator)) { + converter.toVectorSchemaRoot(rows, root); + + WriteParams writeParams = new WriteParams.Builder() + .withMaxRowsPerFile(options.getWriteMaxRowsPerFile()) + .build(); + + List fragments = Fragment.write() + .datasetUri(datasetPath) + .allocator(allocator) + .data(root) + .writeParams(writeParams) + .execute(); + + Overwrite operation = Overwrite.builder().fragments(fragments).schema(arrowSchema).build(); + CommitBuilder builder = new CommitBuilder(datasetPath, allocator) + .writeParams(Collections.emptyMap()); + try (Transaction txn = new Transaction.Builder().operation(operation).build()) { + dataset = builder.execute(txn); + } + } catch (Exception e) { + throw new IOException("Failed to create Lance dataset: " + datasetPath, e); + } + } + + /** + * Apply native upsert via {@code mergeInsert}. + */ + private void mergeInsertRows(List rows) throws IOException { + MergeInsertParams params = new MergeInsertParams(primaryKeys) + .withMatchedUpdateAll() + .withNotMatched(MergeInsertParams.WhenNotMatched.InsertAll); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema, allocator)) { + converter.toVectorSchemaRoot(rows, root); + ArrowArrayStreams.mergeInsert(dataset, params, allocator, root); + } catch (Exception e) { + throw new IOException("Failed to merge-insert Lance rows", e); + } + } + + /** + * Delete rows by an OR-of-AND predicate over the primary-key columns. + */ + private void deleteRows(List rows) { + String predicate = buildDeletePredicate(rows); + LOG.debug("Deleting rows with predicate: {}", predicate); + dataset.delete(predicate); + } + + /** + * Build a SQL predicate of the form {@code (k1 = v1 AND k2 = v2) OR (...)}. + */ + private String buildDeletePredicate(List rows) { + List ors = new ArrayList<>(); + for (RowData row : rows) { + List ands = new ArrayList<>(); + for (int keyIndex : keyIndices) { + String column = rowType.getFieldNames().get(keyIndex); + LogicalType type = rowType.getTypeAt(keyIndex); + Object value = RowDataFieldAccessor.readField(row, keyIndex, type); + ands.add(column + " = " + formatSqlValue(value, type)); + } + ors.add("(" + String.join(" AND ", ands) + ")"); + } + return String.join(" OR ", ors); + } + + /** + * Format a primary-key value as a Lance SQL literal. + */ + private String formatSqlValue(Object value, LogicalType type) { + if (value == null) { + return "NULL"; + } + if (type instanceof TinyIntType || type instanceof SmallIntType + || type instanceof IntType || type instanceof BigIntType + || type instanceof FloatType || type instanceof DoubleType + || type instanceof BooleanType) { + return value.toString(); + } + if (type instanceof VarCharType) { + StringData stringData = (StringData) value; + return "'" + stringData.toString().replace("'", "''") + "'"; + } + throw new UnsupportedOperationException( + "Unsupported primary-key type for delete predicate: " + type.getClass().getSimpleName()); + } + + /** + * Project the primary-key columns into a key that honors equals/hashCode. + */ + private RowData extractKey(RowData value) { + GenericRowData key = new GenericRowData(keyIndices.length); + for (int i = 0; i < keyIndices.length; i++) { + key.setField(i, + RowDataFieldAccessor.readField(value, keyIndices[i], rowType.getTypeAt(keyIndices[i]))); + } + return key; + } + + @Override + public void snapshotState(FunctionSnapshotContext context) throws Exception { + LOG.debug("Snapshot state, checkpointId: {}", context.getCheckpointId()); + flush(); + } + + @Override + public void initializeState(FunctionInitializationContext context) { + LOG.debug("Initialize state, isRestored: {}", context.isRestored()); + } + + @Override + public void close() throws Exception { + LOG.info("Closing Lance Upsert Sink"); + try { + flush(); + } catch (Exception e) { + LOG.warn("Failed to flush data on close", e); + } + if (dataset != null) { + try { + dataset.close(); + } catch (Exception e) { + LOG.warn("Failed to close dataset", e); + } + dataset = null; + } + if (allocator != null) { + try { + allocator.close(); + } catch (Exception e) { + LOG.warn("Failed to close allocator", e); + } + allocator = null; + } + LOG.info("Lance Upsert Sink closed, total written {} rows", totalWrittenRows); + super.close(); + } + + private void deleteDirectory(Path path) throws IOException { + if (Files.isDirectory(path)) { + Files.list(path).forEach(child -> { + try { + deleteDirectory(child); + } catch (IOException e) { + LOG.warn("Failed to delete file: {}", child, e); + } + }); + } + Files.deleteIfExists(path); + } + + public RowType getRowType() { + return rowType; + } + + public List getPrimaryKeys() { + return primaryKeys; + } + + public long getTotalWrittenRows() { + return totalWrittenRows; + } +} diff --git a/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java b/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java new file mode 100644 index 0000000..0311da1 --- /dev/null +++ b/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java @@ -0,0 +1,86 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.lance.Dataset; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * Persists the Flink primary-key column names into the Lance dataset config so that the + * {@code PRIMARY KEY ... NOT ENFORCED} constraint survives a catalog round-trip. + * + *

Lance has no native primary-key constraint; the connector stores the ordered key column + * names under a reserved config key via {@link Dataset#updateConfig(Map)} and restores them via + * {@link Dataset#getConfig()}. + */ +public final class PrimaryKeyPersistence { + + /** Reserved Lance dataset config key holding the comma-separated primary-key columns. */ + public static final String PK_CONFIG_KEY = "flink.primary-key"; + + private PrimaryKeyPersistence() { + // utility class + } + + /** + * Persist the primary-key column names into the dataset config. + * + * @param dataset the Lance dataset (must already be materialized) + * @param primaryKeys ordered primary-key column names; empty/null writes nothing + */ + public static void persist(Dataset dataset, List primaryKeys) { + if (dataset == null || primaryKeys == null || primaryKeys.isEmpty()) { + return; + } + dataset.updateConfig(Collections.singletonMap(PK_CONFIG_KEY, String.join(",", primaryKeys))); + } + + /** + * Load the primary-key column names from the dataset config. + * + * @param dataset the Lance dataset + * @return ordered primary-key column names, or an empty list if not configured + */ + public static List load(Dataset dataset) { + if (dataset == null) { + return Collections.emptyList(); + } + Map config = dataset.getConfig(); + if (config == null) { + return Collections.emptyList(); + } + String raw = config.get(PK_CONFIG_KEY); + if (raw == null || raw.trim().isEmpty()) { + return Collections.emptyList(); + } + List result = new ArrayList<>(); + for (String column : raw.split(",")) { + String trimmed = column.trim(); + if (!trimmed.isEmpty()) { + result.add(trimmed); + } + } + return result; + } +} diff --git a/src/main/java/org/apache/flink/connector/lance/PrimaryKeySelector.java b/src/main/java/org/apache/flink/connector/lance/PrimaryKeySelector.java new file mode 100644 index 0000000..4b5a97d --- /dev/null +++ b/src/main/java/org/apache/flink/connector/lance/PrimaryKeySelector.java @@ -0,0 +1,59 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.LogicalType; + +/** + * Extracts the primary-key projection of a {@link RowData} so that Flink can route all events + * sharing the same key to the same subtask via {@code DataStream#keyBy}. + * + *

The returned key is a freshly allocated {@link GenericRowData} (which implements + * {@code equals}/{@code hashCode}) holding only the primary-key columns, preserving their order. + */ +public class PrimaryKeySelector implements KeySelector { + + private static final long serialVersionUID = 1L; + + private final int[] keyIndices; + private final LogicalType[] keyTypes; + + public PrimaryKeySelector(int[] keyIndices, LogicalType[] keyTypes) { + if (keyIndices == null || keyIndices.length == 0) { + throw new IllegalArgumentException("Primary-key indices must not be empty"); + } + if (keyTypes == null || keyTypes.length != keyIndices.length) { + throw new IllegalArgumentException("Primary-key types must match the key indices"); + } + this.keyIndices = keyIndices; + this.keyTypes = keyTypes; + } + + @Override + public RowData getKey(RowData value) { + GenericRowData key = new GenericRowData(keyIndices.length); + for (int i = 0; i < keyIndices.length; i++) { + key.setField(i, RowDataFieldAccessor.readField(value, keyIndices[i], keyTypes[i])); + } + return key; + } +} diff --git a/src/main/java/org/apache/flink/connector/lance/RowDataFieldAccessor.java b/src/main/java/org/apache/flink/connector/lance/RowDataFieldAccessor.java new file mode 100644 index 0000000..68dac1a --- /dev/null +++ b/src/main/java/org/apache/flink/connector/lance/RowDataFieldAccessor.java @@ -0,0 +1,85 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.BooleanType; +import org.apache.flink.table.types.logical.DoubleType; +import org.apache.flink.table.types.logical.FloatType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.SmallIntType; +import org.apache.flink.table.types.logical.TinyIntType; +import org.apache.flink.table.types.logical.VarCharType; +import org.apache.flink.table.types.logical.VarBinaryType; + +/** + * Type-aware field accessor for {@link RowData}. + * + *

{@link RowData} exposes typed getters rather than a generic {@code getField(int)}, so this + * helper centralizes the dispatch for the scalar types commonly used as primary-key columns. + */ +public final class RowDataFieldAccessor { + + private RowDataFieldAccessor() { + // utility class + } + + /** + * Read a field value honoring nullability. + * + * @return the boxed value, or {@code null} if the field is null + * @throws UnsupportedOperationException for types not supported as a primary-key column + */ + public static Object readField(RowData row, int index, LogicalType type) { + if (row.isNullAt(index)) { + return null; + } + if (type instanceof BooleanType) { + return row.getBoolean(index); + } + if (type instanceof TinyIntType) { + return row.getByte(index); + } + if (type instanceof SmallIntType) { + return row.getShort(index); + } + if (type instanceof IntType) { + return row.getInt(index); + } + if (type instanceof BigIntType) { + return row.getLong(index); + } + if (type instanceof FloatType) { + return row.getFloat(index); + } + if (type instanceof DoubleType) { + return row.getDouble(index); + } + if (type instanceof VarCharType) { + return row.getString(index); + } + if (type instanceof VarBinaryType) { + return row.getBinary(index); + } + throw new UnsupportedOperationException( + "Unsupported field type for primary key: " + type.getClass().getSimpleName()); + } +} diff --git a/src/main/java/org/apache/flink/connector/lance/table/LanceCatalog.java b/src/main/java/org/apache/flink/connector/lance/table/LanceCatalog.java index 600c60a..c4630c4 100644 --- a/src/main/java/org/apache/flink/connector/lance/table/LanceCatalog.java +++ b/src/main/java/org/apache/flink/connector/lance/table/LanceCatalog.java @@ -18,6 +18,7 @@ package org.apache.flink.connector.lance.table; +import org.apache.flink.connector.lance.PrimaryKeyPersistence; import org.apache.flink.connector.lance.converter.LanceTypeConverter; import org.apache.flink.table.api.DataTypes; import org.apache.flink.table.api.Schema; @@ -50,8 +51,11 @@ import org.apache.flink.table.types.logical.RowType; import org.lance.Dataset; +import org.lance.WriteParams; +import org.lance.schema.ColumnAlteration; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.types.pojo.Field; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -103,6 +107,12 @@ public class LanceCatalog extends AbstractCatalog { public static final String DEFAULT_DATABASE = "default"; + /** Option keys that are connector/runtime configuration rather than user TBLPROPERTIES. */ + private static final Set RESERVED_OPTION_KEYS = Collections.unmodifiableSet( + new HashSet<>(java.util.Arrays.asList( + "connector", "path", + "s3-access-key", "s3-secret-key", "s3-region", "s3-endpoint"))); + private final String warehouse; private final Map storageOptions; private final boolean isRemoteStorage; @@ -434,7 +444,13 @@ public CatalogBaseTable getTable(ObjectPath tablePath) throws TableNotExistExcep DataType dataType = LanceTypeConverter.toDataType(field.getType()); schemaBuilder.column(field.getName(), dataType); } - + + // Restore the primary key persisted in the dataset config. + List primaryKeys = PrimaryKeyPersistence.load(dataset); + if (!primaryKeys.isEmpty()) { + schemaBuilder.primaryKey(primaryKeys); + } + Map options = new HashMap<>(); options.put("connector", LanceDynamicTableFactory.IDENTIFIER); options.put("path", datasetPath); @@ -566,16 +582,32 @@ public void createTable(ObjectPath tablePath, CatalogBaseTable table, boolean ig } return; } - + + RowType rowType = resolveRowType(table); + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + LanceTypeConverter.toArrowSchema(rowType); + List primaryKeys = extractPrimaryKeys(table); + + String datasetPath = getDatasetPath(tablePath); + if (isRemoteStorage) { + configureStorageEnvironment(); + } + + // Materialize an empty dataset immediately (matching the community Spark/Trino behavior), + // so the schema and primary-key metadata survive a catalog round-trip before any write. + try (Dataset dataset = Dataset.create( + allocator, datasetPath, arrowSchema, new WriteParams.Builder().build())) { + PrimaryKeyPersistence.persist(dataset, primaryKeys); + } catch (Exception e) { + throw new CatalogException("Failed to create table: " + tablePath, e); + } + if (isRemoteStorage) { - // Remote storage: record table info, actual creation on write String tableKey = tablePath.getDatabaseName() + "/" + tablePath.getObjectName(); knownTables.add(tableKey); } - - // Actual table creation happens on first write - // Only record table metadata here - LOG.info("Registered table: {} (actual dataset will be created on write)", tablePath); + + LOG.info("Created table: {} (primary keys: {})", tablePath, primaryKeys); } @Override @@ -587,9 +619,150 @@ public void alterTable(ObjectPath tablePath, CatalogBaseTable newTable, boolean } return; } - - // Lance does not support modifying table structure - throw new CatalogException("Lance Catalog does not support altering table structure"); + + String datasetPath = getDatasetPath(tablePath); + if (isRemoteStorage) { + configureStorageEnvironment(); + } + + try (Dataset dataset = Dataset.open(datasetPath, allocator)) { + RowType newRowType = resolveRowType(newTable); + RowType oldRowType = LanceTypeConverter.toFlinkRowType(dataset.getSchema()); + + SchemaDiff diff = SchemaDiff.compute(oldRowType, newRowType); + + // Data-type changes are rejected: Lance Java SDK 7.0.0's castTo is verified not to + // mutate the schema (a silent no-op), so supporting it would be a correctness hazard. + if (diff.hasTypeChanges()) { + List affected = new ArrayList<>(diff.getTypeChangedColumns()); + for (SchemaDiff.Rename rename : diff.getRenames()) { + if (rename.getNewType() != null) { + affected.add(rename.getNewName()); + } + } + throw new CatalogException( + "ALTER COLUMN data type change is not supported yet (Lance Java SDK castTo is unreliable). Affected columns: " + + affected); + } + + // Pure renames are metadata-only and reliable. + List alterations = new ArrayList<>(); + for (SchemaDiff.Rename rename : diff.getRenames()) { + alterations.add(new ColumnAlteration.Builder(rename.getOldName()) + .rename(rename.getNewName()) + .build()); + } + if (!alterations.isEmpty()) { + dataset.alterColumns(alterations); + } + + if (!diff.getAddedColumns().isEmpty()) { + List fields = new ArrayList<>(); + for (RowType.RowField added : diff.getAddedColumns()) { + fields.add(LanceTypeConverter.flinkTypeToArrowField(added.getName(), added.getType())); + } + dataset.addColumns(fields); + } + + if (!diff.getDroppedColumns().isEmpty()) { + dataset.dropColumns(diff.getDroppedColumns()); + } + + applyTableProperties(dataset, newTable); + + LOG.info("Altered table: {} (added={}, dropped={})", + tablePath, diff.getAddedColumns().size(), diff.getDroppedColumns().size()); + } catch (CatalogException e) { + throw e; + } catch (Exception e) { + throw new CatalogException("Failed to alter table: " + tablePath, e); + } + } + + /** + * Build a physical {@link RowType} from a table's unresolved schema, skipping computed and + * metadata columns. + */ + private RowType resolveRowType(CatalogBaseTable table) { + List fields = new ArrayList<>(); + for (Schema.UnresolvedColumn column : table.getUnresolvedSchema().getColumns()) { + if (!(column instanceof Schema.UnresolvedPhysicalColumn)) { + continue; + } + Schema.UnresolvedPhysicalColumn physical = (Schema.UnresolvedPhysicalColumn) column; + if (!(physical.getDataType() instanceof DataType)) { + throw new CatalogException( + "Column '" + physical.getName() + "' has an unresolved data type"); + } + DataType dataType = (DataType) physical.getDataType(); + fields.add(new RowType.RowField(physical.getName(), dataType.getLogicalType())); + } + if (fields.isEmpty()) { + throw new CatalogException("Cannot create or alter a Lance table without physical columns"); + } + return new RowType(fields); + } + + /** + * Extract the {@code PRIMARY KEY ... NOT ENFORCED} column names from the table's schema. + * + *

The catalog receives an unresolved schema on the SQL DDL path, whose primary key is + * represented by {@link Schema.UnresolvedPrimaryKey}; the resolved variant + * ({@code UniqueConstraint}) is only available once the schema has been resolved, so this + * deliberately reads the unresolved form. + */ + private List extractPrimaryKeys(CatalogBaseTable table) { + return table.getUnresolvedSchema() + .getPrimaryKey() + .map(Schema.UnresolvedPrimaryKey::getColumnNames) + .orElse(Collections.emptyList()); + } + + /** + * Apply {@code SET}/{@code UNSET} TBLPROPERTIES by syncing non-connector options into the + * Lance dataset config (and removing config keys absent from the new options). + */ + private void applyTableProperties(Dataset dataset, CatalogBaseTable newTable) { + Map options = newTable.getOptions(); + + Map toSet = new HashMap<>(); + for (Map.Entry entry : options.entrySet()) { + if (isTblProperty(entry.getKey())) { + toSet.put(entry.getKey(), entry.getValue()); + } + } + if (!toSet.isEmpty()) { + dataset.updateConfig(toSet); + } + + Set toUnset = new HashSet<>(); + for (String key : dataset.getConfig().keySet()) { + if (PrimaryKeyPersistence.PK_CONFIG_KEY.equals(key)) { + continue; + } + if (isTblProperty(key) && !options.containsKey(key)) { + toUnset.add(key); + } + } + if (!toUnset.isEmpty()) { + dataset.deleteConfigKeys(toUnset); + } + } + + private boolean isTblProperty(String key) { + if (key == null) { + return false; + } + if (RESERVED_OPTION_KEYS.contains(key)) { + return false; + } + if (key.startsWith("hadoop.")) { + return false; + } + return !key.startsWith("read.") + && !key.startsWith("write.") + && !key.startsWith("index.") + && !key.startsWith("vector."); } // ==================== Partition Operations (Lance does not support partitions) ==================== diff --git a/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableFactory.java b/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableFactory.java index 585a6d5..21cb849 100644 --- a/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableFactory.java +++ b/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableFactory.java @@ -22,14 +22,20 @@ import org.apache.flink.configuration.ConfigOptions; import org.apache.flink.configuration.ReadableConfig; import org.apache.flink.connector.lance.config.LanceOptions; +import org.apache.flink.table.catalog.ResolvedSchema; +import org.apache.flink.table.catalog.UniqueConstraint; import org.apache.flink.table.connector.sink.DynamicTableSink; import org.apache.flink.table.connector.source.DynamicTableSource; import org.apache.flink.table.factories.DynamicTableSinkFactory; import org.apache.flink.table.factories.DynamicTableSourceFactory; import org.apache.flink.table.factories.FactoryUtil; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.util.Preconditions; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Set; @@ -210,12 +216,46 @@ public DynamicTableSink createDynamicTableSink(Context context) { ReadableConfig config = helper.getOptions(); LanceOptions options = buildLanceOptions(config, tableOptions); + ResolvedSchema schema = context.getCatalogTable().getResolvedSchema(); + List primaryKeys = extractPrimaryKeys(schema); + int[] primaryKeyIndices = resolvePrimaryKeyIndices(schema, primaryKeys); + return new LanceDynamicTableSink( options, - context.getCatalogTable().getResolvedSchema().toPhysicalRowDataType() + schema.toPhysicalRowDataType(), + primaryKeys, + primaryKeyIndices ); } + /** + * Extract the {@code PRIMARY KEY ... NOT ENFORCED} column names from the resolved schema. + */ + private List extractPrimaryKeys(ResolvedSchema schema) { + return schema.getPrimaryKey() + .map(UniqueConstraint::getColumns) + .orElse(Collections.emptyList()); + } + + /** + * Resolve primary-key column names to physical column indices. + */ + private int[] resolvePrimaryKeyIndices(ResolvedSchema schema, List primaryKeys) { + if (primaryKeys.isEmpty()) { + return new int[0]; + } + RowType rowType = (RowType) schema.toPhysicalRowDataType().getLogicalType(); + List fieldNames = rowType.getFieldNames(); + int[] indices = new int[primaryKeys.size()]; + for (int i = 0; i < primaryKeys.size(); i++) { + int idx = fieldNames.indexOf(primaryKeys.get(i)); + Preconditions.checkArgument( + idx >= 0, "Primary key column '%s' not found in schema", primaryKeys.get(i)); + indices[i] = idx; + } + return indices; + } + /** * 提取以 {@code hadoop.} 为前缀的选项 key，供 {@code validateExcept} 跳过校验。 */ diff --git a/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableSink.java b/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableSink.java index 09ed914..3490200 100644 --- a/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableSink.java +++ b/src/main/java/org/apache/flink/connector/lance/table/LanceDynamicTableSink.java @@ -19,19 +19,26 @@ package org.apache.flink.connector.lance.table; import org.apache.flink.connector.lance.LanceSink; +import org.apache.flink.connector.lance.LanceUpsertSink; +import org.apache.flink.connector.lance.PrimaryKeySelector; import org.apache.flink.connector.lance.config.LanceOptions; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.datastream.DataStreamSink; import org.apache.flink.streaming.api.functions.sink.SinkFunction; import org.apache.flink.table.connector.ChangelogMode; +import org.apache.flink.table.connector.ProviderContext; import org.apache.flink.table.connector.sink.DataStreamSinkProvider; import org.apache.flink.table.connector.sink.DynamicTableSink; import org.apache.flink.table.connector.sink.SinkFunctionProvider; import org.apache.flink.table.data.RowData; import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.LogicalType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.types.RowKind; +import java.util.Collections; +import java.util.List; + /** * Lance dynamic table sink. * @@ -41,17 +48,48 @@ public class LanceDynamicTableSink implements DynamicTableSink { private final LanceOptions options; private final DataType physicalDataType; + private final List primaryKeys; + private final int[] primaryKeyIndices; + private final LogicalType[] primaryKeyTypes; public LanceDynamicTableSink(LanceOptions options, DataType physicalDataType) { + this(options, physicalDataType, Collections.emptyList(), new int[0]); + } + + public LanceDynamicTableSink( + LanceOptions options, + DataType physicalDataType, + List primaryKeys, + int[] primaryKeyIndices) { this.options = options; this.physicalDataType = physicalDataType; + this.primaryKeys = primaryKeys == null ? Collections.emptyList() : primaryKeys; + this.primaryKeyIndices = primaryKeyIndices == null ? new int[0] : primaryKeyIndices; + this.primaryKeyTypes = resolvePrimaryKeyTypes(physicalDataType, this.primaryKeyIndices); + } + + private static LogicalType[] resolvePrimaryKeyTypes(DataType physicalDataType, int[] keyIndices) { + RowType rowType = (RowType) physicalDataType.getLogicalType(); + LogicalType[] types = new LogicalType[keyIndices.length]; + for (int i = 0; i < keyIndices.length; i++) { + types[i] = rowType.getTypeAt(keyIndices[i]); + } + return types; } @Override public ChangelogMode getChangelogMode(ChangelogMode requestedMode) { - // Lance only supports INSERT operations + if (primaryKeys.isEmpty()) { + // No primary key: insert-only (append). + return ChangelogMode.newBuilder() + .addContainedKind(RowKind.INSERT) + .build(); + } + // With a primary key: support upsert (+I/+U) and delete (-D). return ChangelogMode.newBuilder() .addContainedKind(RowKind.INSERT) + .addContainedKind(RowKind.UPDATE_AFTER) + .addContainedKind(RowKind.DELETE) .build(); } @@ -59,15 +97,29 @@ public ChangelogMode getChangelogMode(ChangelogMode requestedMode) { public SinkRuntimeProvider getSinkRuntimeProvider(Context context) { RowType rowType = (RowType) physicalDataType.getLogicalType(); - // Create LanceSink - LanceSink lanceSink = new LanceSink(options, rowType); + if (primaryKeys.isEmpty()) { + // Append-only path: keep the existing LanceSink. + LanceSink lanceSink = new LanceSink(options, rowType); + return SinkFunctionProvider.of(lanceSink); + } - return SinkFunctionProvider.of(lanceSink); + // Keyed upsert path: keyBy the primary key so all events for a key reach one subtask. + return new DataStreamSinkProvider() { + @Override + public DataStreamSink consumeDataStream( + ProviderContext providerContext, DataStream dataStream) { + DataStream keyed = + dataStream.keyBy(new PrimaryKeySelector(primaryKeyIndices, primaryKeyTypes)); + LanceUpsertSink upsertSink = + new LanceUpsertSink(options, rowType, primaryKeys, primaryKeyIndices); + return keyed.addSink(upsertSink); + } + }; } @Override public DynamicTableSink copy() { - return new LanceDynamicTableSink(options, physicalDataType); + return new LanceDynamicTableSink(options, physicalDataType, primaryKeys, primaryKeyIndices); } @Override diff --git a/src/main/java/org/apache/flink/connector/lance/table/SchemaDiff.java b/src/main/java/org/apache/flink/connector/lance/table/SchemaDiff.java new file mode 100644 index 0000000..9a01af9 --- /dev/null +++ b/src/main/java/org/apache/flink/connector/lance/table/SchemaDiff.java @@ -0,0 +1,214 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance.table; + +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Column-level diff between an existing Lance dataset schema (old) and the schema requested by an + * {@code ALTER TABLE} statement (new). + * + *

Detects {@code ADD COLUMN}, {@code DROP COLUMN}, in-place type changes and column renames. + * A rename is inferred when the old and new schemas differ by exactly one removed column and one + * added column at the same position — Flink never issues {@code ADD} and {@code DROP} in a single + * {@code ALTER TABLE}, so such a pairing can only be a {@code RENAME COLUMN}. If the diff cannot + * be safely classified (an unsafe drop+add that would lose data), {@link #compute} throws rather + * than guessing. + */ +public final class SchemaDiff { + + private final List addedColumns; + private final List droppedColumns; + private final List typeChangedColumns; + private final Map typeChangeByColumn; + private final List renames; + + private SchemaDiff( + List addedColumns, + List droppedColumns, + Map typeChangeByColumn, + List renames) { + this.addedColumns = Collections.unmodifiableList(addedColumns); + this.droppedColumns = Collections.unmodifiableList(droppedColumns); + this.typeChangeByColumn = Collections.unmodifiableMap(typeChangeByColumn); + this.typeChangedColumns = + Collections.unmodifiableList(new ArrayList<>(typeChangeByColumn.keySet())); + this.renames = Collections.unmodifiableList(renames); + } + + public static SchemaDiff compute(RowType oldRowType, RowType newRowType) { + Map oldFields = new LinkedHashMap<>(); + for (RowType.RowField field : oldRowType.getFields()) { + oldFields.put(field.getName(), field); + } + + Map newFields = new LinkedHashMap<>(); + for (RowType.RowField field : newRowType.getFields()) { + newFields.put(field.getName(), field); + } + + Map typeChanges = new LinkedHashMap<>(); + List droppedCandidates = new ArrayList<>(); + List addedCandidates = new ArrayList<>(); + + for (RowType.RowField newField : newRowType.getFields()) { + RowType.RowField oldField = oldFields.get(newField.getName()); + if (oldField == null) { + addedCandidates.add(newField); + } else if (!oldField.getType().equals(newField.getType())) { + typeChanges.put(newField.getName(), newField.getType()); + } + } + for (RowType.RowField oldField : oldRowType.getFields()) { + if (!newFields.containsKey(oldField.getName())) { + droppedCandidates.add(oldField); + } + } + + List added = new ArrayList<>(); + List dropped = new ArrayList<>(); + List renames = new ArrayList<>(); + + if (!droppedCandidates.isEmpty() && !addedCandidates.isEmpty()) { + // A single ALTER never combines ADD and DROP; a paired remove+add at the same position + // can only be RENAME COLUMN. Anything else is unsafe and must be rejected to avoid + // silently dropping the old column's data. + if (droppedCandidates.size() != addedCandidates.size()) { + throw new IllegalArgumentException( + "Cannot safely distinguish RENAME from DROP+ADD: dropped " + + fieldNames(droppedCandidates) + " but added " + + fieldNames(addedCandidates)); + } + + List oldNames = oldRowType.getFieldNames(); + List newNames = newRowType.getFieldNames(); + for (int i = 0; i < droppedCandidates.size(); i++) { + RowType.RowField oldField = droppedCandidates.get(i); + RowType.RowField newField = addedCandidates.get(i); + if (oldNames.indexOf(oldField.getName()) != newNames.indexOf(newField.getName())) { + throw new IllegalArgumentException( + "Cannot safely distinguish RENAME from DROP+ADD: column positions differ for '" + + oldField.getName() + "' -> '" + newField.getName() + "'"); + } + LogicalType newType = oldField.getType().equals(newField.getType()) + ? null : newField.getType(); + renames.add(new Rename(oldField.getName(), newField.getName(), newType)); + } + } else { + added = addedCandidates; + for (RowType.RowField oldField : droppedCandidates) { + dropped.add(oldField.getName()); + } + } + + return new SchemaDiff(added, dropped, typeChanges, renames); + } + + private static List fieldNames(List fields) { + List names = new ArrayList<>(); + for (RowType.RowField field : fields) { + names.add(field.getName()); + } + return names; + } + + public List getAddedColumns() { + return addedColumns; + } + + public List getDroppedColumns() { + return droppedColumns; + } + + /** + * Columns whose data type changed in place (same name, different type). + */ + public List getTypeChangedColumns() { + return typeChangedColumns; + } + + /** + * In-place type changes keyed by column name, mapping to the new type. + */ + public Map getTypeChangeByColumn() { + return typeChangeByColumn; + } + + /** + * Column renames detected as a remove+add pair at the same position. + */ + public List getRenames() { + return renames; + } + + public boolean hasTypeChanges() { + if (!typeChangedColumns.isEmpty()) { + return true; + } + for (Rename rename : renames) { + if (rename.getNewType() != null) { + return true; + } + } + return false; + } + + public boolean isEmpty() { + return addedColumns.isEmpty() + && droppedColumns.isEmpty() + && typeChangeByColumn.isEmpty() + && renames.isEmpty(); + } + + /** + * A detected {@code RENAME COLUMN}, optionally combined with a data type change. + */ + public static final class Rename { + + private final String oldName; + private final String newName; + private final LogicalType newType; + + Rename(String oldName, String newName, LogicalType newType) { + this.oldName = oldName; + this.newName = newName; + this.newType = newType; + } + + public String getOldName() { + return oldName; + } + + public String getNewName() { + return newName; + } + + /** The new type when the rename also changes the type, otherwise {@code null}. */ + public LogicalType getNewType() { + return newType; + } + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/CompositePkDeleteITCase.java b/src/test/java/org/apache/flink/connector/lance/CompositePkDeleteITCase.java new file mode 100644 index 0000000..0fe23d1 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/CompositePkDeleteITCase.java @@ -0,0 +1,173 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.lance.CommitBuilder; +import org.lance.Dataset; +import org.lance.Fragment; +import org.lance.FragmentMetadata; +import org.lance.Transaction; +import org.lance.WriteParams; +import org.lance.operation.Overwrite; + +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.fail; + +/** + * Regression test for composite-primary-key DELETE predicate syntax (issue #63, Phase 0 Q3). + * + *

Locks in the spike finding that Lance's filter parser does not support row-value + * tuple {@code IN} for composite keys, and that the OR-of-AND shape is the only working form. + * This is the acceptance basis for Phase B's predicate builder. + * + *

Verified against {@code org.lance:lance-core:7.0.0} (arrow 18.3.0, JDK 11): + * + *

    + *
  • {@code (user_id, event_ts) IN ((..), (..))} → rejected with + * {@code IllegalArgumentException: Expression '(user_id, event_ts)' is not supported SQL}.
    • + *
  • {@code (user_id = .. AND event_ts = timestamp '..') OR (..)} → accepted and deletes + * exactly the matching rows.
    • + *
+ */ +class CompositePkDeleteITCase { + + @TempDir + Path tempDir; + + @BeforeAll + static void ensureArrowNettyLoaded() { + // Force Arrow to use the Netty allocator, matching LanceTimeTravelITCase and + // LanceNamespaceCatalogITCase to avoid classloader-related SPI issues. + System.setProperty("arrow.memory.allocator.type", "Netty"); + try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { + // allocator created and closed successfully + } + } + + // 2023-11-14T22:13:20Z / 22:13:21Z / 22:13:22Z in epoch microseconds. + private static final long TS1 = 1_700_000_000_000_000L; + private static final long TS2 = 1_700_000_001_000_000L; + private static final long TS3 = 1_700_000_002_000_000L; + + private static Schema compositeSchema() { + return new Schema(Arrays.asList( + new Field("user_id", FieldType.nullable(new ArrowType.Int(64, true)), null), + new Field("event_ts", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)), null), + new Field("payload", FieldType.nullable(ArrowType.Utf8.INSTANCE), null) + )); + } + + @Test + @DisplayName("tuple-IN is rejected; OR-of-AND deletes exactly the matching composite-PK rows") + void compositePkDeleteMustUseOrOfAnd() throws Exception { + String tupleInUri = tempDir.resolve("tuple_in").toString(); + seedComposite(tupleInUri); + + // tuple-IN must be rejected by the Lance filter parser. + String tupleIn = + "(user_id, event_ts) IN ((100, timestamp '2023-11-14 22:13:20')," + + " (200, timestamp '2023-11-14 22:13:22'))"; + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(tupleInUri, allocator)) { + try { + ds.delete(tupleIn); + fail("tuple-IN predicate must be rejected by the Lance filter parser"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("not supported"); + } + } + + // OR-of-AND must succeed and remove exactly the two matched rows (TS1 and TS3), + // leaving only the TS2 row. + String orOfAndUri = tempDir.resolve("or_of_and").toString(); + seedComposite(orOfAndUri); + String orOfAnd = + "(user_id = 100 AND event_ts = timestamp '2023-11-14 22:13:20')" + + " OR " + + "(user_id = 200 AND event_ts = timestamp '2023-11-14 22:13:22')"; + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(orOfAndUri, allocator)) { + ds.delete(orOfAnd); + + assertThat(ds.countRows()).isEqualTo(1L); + // The surviving row is (user_id=100, event_ts=TS2, payload="b"). + assertThat(ds.countRows("user_id = 100")).isEqualTo(1L); + assertThat(ds.countRows("user_id = 200")).isZero(); + assertThat(ds.countRows("event_ts = timestamp '2023-11-14 22:13:21'")).isEqualTo(1L); + } + } + + private void seedComposite(String uri) throws Exception { + Schema schema = compositeSchema(); + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + BigIntVector userVec = (BigIntVector) root.getVector("user_id"); + TimeStampMicroVector tsVec = (TimeStampMicroVector) root.getVector("event_ts"); + VarCharVector payloadVec = (VarCharVector) root.getVector("payload"); + root.setRowCount(3); + userVec.setSafe(0, 100L); + tsVec.setSafe(0, TS1); + payloadVec.setSafe(0, "a".getBytes(StandardCharsets.UTF_8)); + userVec.setSafe(1, 100L); + tsVec.setSafe(1, TS2); + payloadVec.setSafe(1, "b".getBytes(StandardCharsets.UTF_8)); + userVec.setSafe(2, 200L); + tsVec.setSafe(2, TS3); + payloadVec.setSafe(2, "c".getBytes(StandardCharsets.UTF_8)); + + WriteParams writeParams = new WriteParams.Builder().withMaxRowsPerFile(1_000_000).build(); + List fragments = Fragment.write() + .datasetUri(uri) + .allocator(allocator) + .data(root) + .writeParams(writeParams) + .execute(); + + CommitBuilder builder = new CommitBuilder(uri, allocator).writeParams(Collections.emptyMap()); + try (Transaction txn = new Transaction.Builder() + .operation(Overwrite.builder().fragments(fragments).schema(schema).build()) + .build(); + Dataset ds = builder.execute(txn)) { + assertThat(ds.countRows()).isEqualTo(3L); + } + } + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/LanceSchemaEvolutionITCase.java b/src/test/java/org/apache/flink/connector/lance/LanceSchemaEvolutionITCase.java new file mode 100644 index 0000000..c9fd1be --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/LanceSchemaEvolutionITCase.java @@ -0,0 +1,185 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.flink.connector.lance.converter.LanceTypeConverter; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.VarCharType; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.lance.CommitBuilder; +import org.lance.Dataset; +import org.lance.Fragment; +import org.lance.FragmentMetadata; +import org.lance.Transaction; +import org.lance.WriteParams; +import org.lance.operation.Overwrite; +import org.lance.schema.ColumnAlteration; + +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Verifies the Lance SDK {@code addColumns}/{@code dropColumns} behavior that backs + * {@code ALTER TABLE ADD/DROP COLUMN} in {@code LanceCatalog#alterTable}. + */ +class LanceSchemaEvolutionITCase { + + @TempDir + Path tempDir; + + @BeforeAll + static void ensureArrowNettyLoaded() { + System.setProperty("arrow.memory.allocator.type", "Netty"); + try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { + // allocator created and closed successfully + } + } + + private void seed(String uri) throws Exception { + Schema schema = new Schema(Collections.singletonList( + new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), null))); + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + BigIntVector idVec = (BigIntVector) root.getVector("id"); + root.setRowCount(1); + idVec.setSafe(0, 1L); + + WriteParams writeParams = new WriteParams.Builder().withMaxRowsPerFile(1_000_000).build(); + List fragments = Fragment.write() + .datasetUri(uri) + .allocator(allocator) + .data(root) + .writeParams(writeParams) + .execute(); + + CommitBuilder builder = new CommitBuilder(uri, allocator).writeParams(Collections.emptyMap()); + try (Transaction txn = new Transaction.Builder() + .operation(Overwrite.builder().fragments(fragments).schema(schema).build()) + .build(); + Dataset ds = builder.execute(txn)) { + assertThat(ds.countRows()).isEqualTo(1L); + } + } + } + + @Test + @DisplayName("addColumns then dropColumns mutates the dataset schema") + void addAndDropColumns() throws Exception { + String uri = tempDir.resolve("schema_evo").toString(); + seed(uri); + + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(uri, allocator)) { + // Add a 'name' column. + Field nameField = LanceTypeConverter.flinkTypeToArrowField( + "name", new VarCharType(true, VarCharType.MAX_LENGTH)); + ds.addColumns(Collections.singletonList(nameField)); + + List columns = ds.getSchema().getFields().stream() + .map(Field::getName) + .collect(Collectors.toList()); + assertThat(columns).containsExactlyInAnyOrder("id", "name"); + + // Drop the 'name' column. + ds.dropColumns(Collections.singletonList("name")); + List afterDrop = ds.getSchema().getFields().stream() + .map(Field::getName) + .collect(Collectors.toList()); + assertThat(afterDrop).containsExactly("id"); + } + } + + @Test + @DisplayName("type change detection helper flags a type change") + void typeChangeDetection() { + Field old = new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), null); + Field changed = new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null); + assertThat(old.getType().equals(changed.getType())).isFalse(); + } + + private void seedWithName(String uri) throws Exception { + Schema schema = new Schema(Arrays.asList( + new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("name", FieldType.nullable(ArrowType.Utf8.INSTANCE), null))); + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + IntVector idVec = (IntVector) root.getVector("id"); + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + root.setRowCount(1); + idVec.setSafe(0, 42); + nameVec.setSafe(0, "alice".getBytes(StandardCharsets.UTF_8)); + + WriteParams writeParams = new WriteParams.Builder().withMaxRowsPerFile(1_000_000).build(); + List fragments = Fragment.write() + .datasetUri(uri) + .allocator(allocator) + .data(root) + .writeParams(writeParams) + .execute(); + + CommitBuilder builder = new CommitBuilder(uri, allocator).writeParams(Collections.emptyMap()); + try (Transaction txn = new Transaction.Builder() + .operation(Overwrite.builder().fragments(fragments).schema(schema).build()) + .build(); + Dataset ds = builder.execute(txn)) { + assertThat(ds.countRows()).isEqualTo(1L); + } + } + } + + @Test + @DisplayName("alterColumns renames a column and preserves its data") + void renameColumn() throws Exception { + String uri = tempDir.resolve("rename").toString(); + seedWithName(uri); + + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(uri, allocator)) { + ds.alterColumns(Collections.singletonList( + new ColumnAlteration.Builder("name").rename("full_name").build())); + + List columns = ds.getSchema().getFields().stream() + .map(Field::getName) + .collect(Collectors.toList()); + assertThat(columns).containsExactlyInAnyOrder("id", "full_name"); + assertThat(ds.countRows()).isEqualTo(1L); + assertThat(ds.countRows("full_name = 'alice'")).isEqualTo(1L); + } + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/LanceSinkConcurrencyITCase.java b/src/test/java/org/apache/flink/connector/lance/LanceSinkConcurrencyITCase.java new file mode 100644 index 0000000..aa7608a --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/LanceSinkConcurrencyITCase.java @@ -0,0 +1,113 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.connector.lance.config.LanceOptions; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.lance.Dataset; + +import java.nio.file.Path; +import java.util.Arrays; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration test for append-mode {@link LanceSink} concurrency. + * + *

Locks in that two subtasks concurrently first-writing to a not-yet-existing dataset do not + * clobber each other (regression: {@code Overwrite} on first write silently dropped the other + * subtask's rows). + */ +class LanceSinkConcurrencyITCase { + + @TempDir + Path tempDir; + + @BeforeAll + static void ensureArrowNettyLoaded() { + System.setProperty("arrow.memory.allocator.type", "Netty"); + try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { + // allocator created and closed successfully + } + } + + private static RowType rowType() { + return new RowType(Arrays.asList( + new RowType.RowField("id", new BigIntType(false)), + new RowType.RowField("content", new VarCharType(true, VarCharType.MAX_LENGTH)))); + } + + private LanceOptions options(String path) { + return LanceOptions.builder() + .path(path) + .writeBatchSize(100) + .writeMode(LanceOptions.WriteMode.APPEND) + .build(); + } + + private GenericRowData row(long id, String content) { + GenericRowData row = new GenericRowData(2); + row.setField(0, id); + row.setField(1, StringData.fromString(content)); + return row; + } + + private long countRows(String path) { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(path, allocator)) { + return ds.countRows(); + } + } + + @Test + @DisplayName("two append subtasks concurrently first-write distinct rows without data loss") + void twoSubtasksConcurrentFirstWrite() throws Exception { + String path = tempDir.resolve("concurrent_append").toString(); + + LanceSink s1 = new LanceSink(options(path), rowType()); + LanceSink s2 = new LanceSink(options(path), rowType()); + s1.open(new Configuration()); + s2.open(new Configuration()); + + try { + s1.invoke(row(1L, "one"), null); + s1.flush(); + s2.invoke(row(2L, "two"), null); + s2.flush(); + } finally { + s1.close(); + s2.close(); + } + + // The second first-write must not clobber the first (regression: Overwrite clobbering). + assertThat(countRows(path)).isEqualTo(2L); + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkITCase.java b/src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkITCase.java new file mode 100644 index 0000000..e282040 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkITCase.java @@ -0,0 +1,314 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.connector.lance.config.LanceOptions; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; +import org.apache.flink.types.RowKind; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.lance.Dataset; + +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration test for {@link LanceUpsertSink}: primary-key upsert (+I/+U) via {@code mergeInsert} + * and delete (-D) via {@code Dataset#delete}. + */ +class LanceUpsertSinkITCase { + + @TempDir + Path tempDir; + + @BeforeAll + static void ensureArrowNettyLoaded() { + System.setProperty("arrow.memory.allocator.type", "Netty"); + try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { + // allocator created and closed successfully + } + } + + private static RowType rowType() { + return new RowType(Arrays.asList( + new RowType.RowField("id", new BigIntType(false)), + new RowType.RowField("name", new VarCharType(true, VarCharType.MAX_LENGTH)))); + } + + private static final String[] PRIMARY_KEYS = {"id"}; + private static final int[] KEY_INDICES = {0}; + + private LanceOptions options(String path) { + return LanceOptions.builder() + .path(path) + .writeBatchSize(100) + .writeMode(LanceOptions.WriteMode.APPEND) + .build(); + } + + private GenericRowData row(long id, String name, RowKind kind) { + GenericRowData row = new GenericRowData(2); + row.setField(0, id); + row.setField(1, StringData.fromString(name)); + row.setRowKind(kind); + return row; + } + + private long countRows(String path) { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(path, allocator)) { + return ds.countRows(); + } + } + + @Test + @DisplayName("first flush creates the dataset and persists the primary key") + void firstFlushCreatesDatasetAndPersistsPrimaryKey() throws Exception { + String path = tempDir.resolve("first").toString(); + LanceUpsertSink sink = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + + sink.open(new Configuration()); + try { + sink.invoke(row(1L, "alice", RowKind.INSERT), null); + sink.invoke(row(2L, "bob", RowKind.INSERT), null); + sink.flush(); + } finally { + sink.close(); + } + + assertThat(countRows(path)).isEqualTo(2L); + + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(path, allocator)) { + assertThat(PrimaryKeyPersistence.load(ds)).containsExactly("id"); + } + } + + @Test + @DisplayName("+I(k) then -D(k) in the same flush leaves k absent") + void insertThenDeleteSameFlushLeavesKeyAbsent() throws Exception { + String path = tempDir.resolve("insert_then_delete").toString(); + LanceUpsertSink sink = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + + sink.open(new Configuration()); + try { + // Seed one row in a first flush. + sink.invoke(row(1L, "alice", RowKind.INSERT), null); + sink.flush(); + assertThat(countRows(path)).isEqualTo(1L); + + // +I(1) then -D(1) collapse to delete. + sink.invoke(row(1L, "alice", RowKind.INSERT), null); + sink.invoke(row(1L, "alice", RowKind.DELETE), null); + sink.flush(); + } finally { + sink.close(); + } + + assertThat(countRows(path)).isZero(); + } + + @Test + @DisplayName("-D(k) then +I(k) in the same flush leaves k present") + void deleteThenInsertSameFlushLeavesKeyPresent() throws Exception { + String path = tempDir.resolve("delete_then_insert").toString(); + LanceUpsertSink sink = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + + sink.open(new Configuration()); + try { + sink.invoke(row(1L, "alice", RowKind.INSERT), null); + sink.flush(); + assertThat(countRows(path)).isEqualTo(1L); + + // -D(1) then +I(1) collapse to upsert of the new value. + sink.invoke(row(1L, "alice", RowKind.DELETE), null); + sink.invoke(row(1L, "alice-new", RowKind.INSERT), null); + sink.flush(); + } finally { + sink.close(); + } + + assertThat(countRows(path)).isEqualTo(1L); + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(path, allocator)) { + assertThat(ds.countRows("name = 'alice-new'")).isEqualTo(1L); + } + } + + @Test + @DisplayName("+U(k) replaces the previous value via merge-insert") + void updateAfterReplacesPreviousValue() throws Exception { + String path = tempDir.resolve("update").toString(); + LanceUpsertSink sink = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + + sink.open(new Configuration()); + try { + sink.invoke(row(1L, "old", RowKind.INSERT), null); + sink.flush(); + + sink.invoke(row(1L, "new", RowKind.UPDATE_AFTER), null); + sink.flush(); + } finally { + sink.close(); + } + + assertThat(countRows(path)).isEqualTo(1L); + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset ds = Dataset.open(path, allocator)) { + assertThat(ds.countRows("name = 'new'")).isEqualTo(1L); + assertThat(ds.countRows("name = 'old'")).isZero(); + } + } + + @Test + @DisplayName("replaying the same upsert batch (checkpoint replay) is idempotent") + void replaySameUpsertBatchIsIdempotent() throws Exception { + String path = tempDir.resolve("replay_upsert").toString(); + LanceUpsertSink sink = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + + sink.open(new Configuration()); + try { + // First flush (checkpoint 1 completed). + sink.invoke(row(1L, "alice", RowKind.INSERT), null); + sink.invoke(row(2L, "bob", RowKind.INSERT), null); + sink.flush(); + assertThat(countRows(path)).isEqualTo(2L); + + // Replay the same batch (checkpoint 1 failed and was retried). + sink.invoke(row(1L, "alice", RowKind.INSERT), null); + sink.invoke(row(2L, "bob", RowKind.INSERT), null); + sink.flush(); + } finally { + sink.close(); + } + + // Idempotent: replay must not duplicate rows. + assertThat(countRows(path)).isEqualTo(2L); + } + + @Test + @DisplayName("replaying a delete batch (checkpoint replay) is idempotent") + void replayDeleteBatchIsIdempotent() throws Exception { + String path = tempDir.resolve("replay_delete").toString(); + LanceUpsertSink sink = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + + sink.open(new Configuration()); + try { + sink.invoke(row(1L, "alice", RowKind.INSERT), null); + sink.flush(); + assertThat(countRows(path)).isEqualTo(1L); + + // First delete flush. + sink.invoke(row(1L, "alice", RowKind.DELETE), null); + sink.flush(); + assertThat(countRows(path)).isZero(); + + // Replay the delete (checkpoint 2 failed and was retried). + sink.invoke(row(1L, "alice", RowKind.DELETE), null); + sink.flush(); + } finally { + sink.close(); + } + + // Idempotent: replaying a delete of an absent key must remain a no-op. + assertThat(countRows(path)).isZero(); + } + + @Test + @DisplayName("two subtasks concurrently write distinct keys to an existing dataset") + void twoSubtasksConcurrentWriteToExistingDataset() throws Exception { + String path = tempDir.resolve("concurrent_existing").toString(); + + // Seed the dataset with key=1 so it already exists. + LanceUpsertSink seeder = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + seeder.open(new Configuration()); + seeder.invoke(row(1L, "one", RowKind.INSERT), null); + seeder.flush(); + seeder.close(); + + // Two subtasks open at the same base version. + LanceUpsertSink s1 = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + LanceUpsertSink s2 = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + s1.open(new Configuration()); + s2.open(new Configuration()); + + try { + s1.invoke(row(2L, "two", RowKind.INSERT), null); + s1.flush(); + s2.invoke(row(3L, "three", RowKind.INSERT), null); + s2.flush(); + } finally { + s1.close(); + s2.close(); + } + + // Both writes must land: no write conflict, no lost row. + assertThat(countRows(path)).isEqualTo(3L); + } + + @Test + @DisplayName("two subtasks concurrently first-write distinct keys without data loss") + void twoSubtasksConcurrentFirstWrite() throws Exception { + String path = tempDir.resolve("concurrent_first").toString(); + + // Dataset does NOT exist yet. Both subtasks open, then first-write. + LanceUpsertSink s1 = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + LanceUpsertSink s2 = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + s1.open(new Configuration()); + s2.open(new Configuration()); + + try { + s1.invoke(row(1L, "one", RowKind.INSERT), null); + s1.flush(); + s2.invoke(row(2L, "two", RowKind.INSERT), null); + s2.flush(); + } finally { + s1.close(); + s2.close(); + } + + // The second first-write must not clobber the first (regression: Overwrite clobbering). + assertThat(countRows(path)).isEqualTo(2L); + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/table/LanceCatalogTableITCase.java b/src/test/java/org/apache/flink/connector/lance/table/LanceCatalogTableITCase.java new file mode 100644 index 0000000..bcb07ba --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/table/LanceCatalogTableITCase.java @@ -0,0 +1,196 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance.table; + +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.catalog.CatalogTable; +import org.apache.flink.table.catalog.ObjectPath; +import org.apache.flink.table.catalog.exceptions.CatalogException; + +import org.apache.arrow.memory.RootAllocator; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Path; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Integration test for the directory-backed {@link LanceCatalog} table lifecycle. + * + *

Locks in that {@code createTable} materializes an empty Lance dataset immediately (rather + * than deferring to first write), so the schema and primary-key constraint survive a catalog + * round-trip before any data is written — matching the community Spark/Trino behavior. + */ +class LanceCatalogTableITCase { + + @TempDir + Path tempDir; + + private LanceCatalog catalog; + + @BeforeAll + static void ensureArrowNettyLoaded() { + System.setProperty("arrow.memory.allocator.type", "Netty"); + try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { + // allocator created and closed successfully + } + } + + @BeforeEach + void setUp() { + String warehouse = tempDir.resolve("warehouse").toString(); + catalog = new LanceCatalog("test", "default", warehouse); + catalog.open(); + } + + @AfterEach + void tearDown() { + catalog.close(); + } + + @Test + @DisplayName("createTable materializes an empty dataset immediately") + void createTableMaterializesEmptyDatasetImmediately() throws Exception { + catalog.createDatabase("db", null, false); + ObjectPath tablePath = new ObjectPath("db", "t"); + + catalog.createTable(tablePath, CatalogTable.of( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), false); + + // The empty dataset must be materialized on disk right away, not on first write. + assertThat(catalog.tableExists(tablePath)).isTrue(); + assertThat(catalog.listTables("db")).contains("t"); + } + + @Test + @DisplayName("createTable persists and restores the primary key") + void createTablePersistsAndRestoresPrimaryKey() throws Exception { + catalog.createDatabase("db", null, false); + ObjectPath tablePath = new ObjectPath("db", "t"); + + catalog.createTable(tablePath, CatalogTable.of( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .primaryKey("id") + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), false); + + // The primary key must survive a round-trip via the dataset config. + CatalogTable loaded = (CatalogTable) catalog.getTable(tablePath); + Schema unresolved = loaded.getUnresolvedSchema(); + assertThat(unresolved.getPrimaryKey()).isPresent(); + assertThat(unresolved.getPrimaryKey().get().getColumnNames()).containsExactly("id"); + } + + @Test + @DisplayName("createTable without a primary key restores no constraint") + void createTableWithoutPrimaryKeyRestoresNoConstraint() throws Exception { + catalog.createDatabase("db", null, false); + ObjectPath tablePath = new ObjectPath("db", "t"); + + catalog.createTable(tablePath, CatalogTable.of( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), false); + + CatalogTable loaded = (CatalogTable) catalog.getTable(tablePath); + assertThat(loaded.getUnresolvedSchema().getPrimaryKey()).isEmpty(); + } + + @Test + @DisplayName("alterTable renames a column") + void alterTableRenamesColumn() throws Exception { + catalog.createDatabase("db", null, false); + ObjectPath tablePath = new ObjectPath("db", "t"); + + catalog.createTable(tablePath, CatalogTable.of( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), false); + + catalog.alterTable(tablePath, CatalogTable.of( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("full_name", DataTypes.STRING()) + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), false); + + CatalogTable loaded = (CatalogTable) catalog.getTable(tablePath); + List columns = loaded.getUnresolvedSchema().getColumns().stream() + .map(c -> ((Schema.UnresolvedPhysicalColumn) c).getName()) + .collect(Collectors.toList()); + assertThat(columns).containsExactlyInAnyOrder("id", "full_name"); + } + + @Test + @DisplayName("alterTable rejects a data type change") + void alterTableRejectsTypeChange() throws Exception { + catalog.createDatabase("db", null, false); + ObjectPath tablePath = new ObjectPath("db", "t"); + + catalog.createTable(tablePath, CatalogTable.of( + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column("name", DataTypes.STRING()) + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), false); + + assertThatThrownBy(() -> catalog.alterTable(tablePath, CatalogTable.of( + Schema.newBuilder() + .column("id", DataTypes.BIGINT()) + .column("name", DataTypes.STRING()) + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), false)) + .isInstanceOf(CatalogException.class) + .hasMessageContaining("data type change"); + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/table/SchemaDiffTest.java b/src/test/java/org/apache/flink/connector/lance/table/SchemaDiffTest.java new file mode 100644 index 0000000..b7bef98 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/table/SchemaDiffTest.java @@ -0,0 +1,146 @@ +/* + * 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. + */ + +package org.apache.flink.connector.lance.table; + +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class SchemaDiffTest { + + private static RowType rowType(RowType.RowField... fields) { + return new RowType(Arrays.asList(fields)); + } + + private static final RowType.RowField ID = new RowType.RowField("id", new BigIntType(false)); + private static final RowType.RowField NAME = + new RowType.RowField("name", new VarCharType(true, VarCharType.MAX_LENGTH)); + private static final RowType.RowField AGE = new RowType.RowField("age", new BigIntType(true)); + + @Test + void detectsAddedColumn() { + RowType oldRow = rowType(ID); + RowType newRow = rowType(ID, NAME); + + SchemaDiff diff = SchemaDiff.compute(oldRow, newRow); + + assertThat(diff.getAddedColumns()).extracting(RowType.RowField::getName).containsExactly("name"); + assertThat(diff.getDroppedColumns()).isEmpty(); + assertThat(diff.hasTypeChanges()).isFalse(); + } + + @Test + void detectsDroppedColumn() { + RowType oldRow = rowType(ID, NAME); + RowType newRow = rowType(ID); + + SchemaDiff diff = SchemaDiff.compute(oldRow, newRow); + + assertThat(diff.getAddedColumns()).isEmpty(); + assertThat(diff.getDroppedColumns()).containsExactly("name"); + assertThat(diff.hasTypeChanges()).isFalse(); + } + + @Test + void detectsTypeChange() { + RowType oldRow = rowType(ID, new RowType.RowField("name", new VarCharType(true, VarCharType.MAX_LENGTH))); + RowType newRow = rowType(ID, new RowType.RowField("name", new BigIntType(true))); + + SchemaDiff diff = SchemaDiff.compute(oldRow, newRow); + + assertThat(diff.getTypeChangedColumns()).containsExactly("name"); + assertThat(diff.hasTypeChanges()).isTrue(); + } + + @Test + void reportsEmptyDiffForIdenticalSchemas() { + RowType row = rowType(ID, NAME); + + SchemaDiff diff = SchemaDiff.compute(row, row); + + assertThat(diff.isEmpty()).isTrue(); + assertThat(diff.getAddedColumns()).isEmpty(); + assertThat(diff.getDroppedColumns()).isEmpty(); + assertThat(diff.hasTypeChanges()).isFalse(); + } + + @Test + void detectsRename() { + RowType oldRow = rowType(ID, NAME); + RowType newRow = rowType(ID, + new RowType.RowField("full_name", new VarCharType(true, VarCharType.MAX_LENGTH))); + + SchemaDiff diff = SchemaDiff.compute(oldRow, newRow); + + assertThat(diff.getRenames()).hasSize(1); + SchemaDiff.Rename rename = diff.getRenames().get(0); + assertThat(rename.getOldName()).isEqualTo("name"); + assertThat(rename.getNewName()).isEqualTo("full_name"); + assertThat(rename.getNewType()).isNull(); + assertThat(diff.getAddedColumns()).isEmpty(); + assertThat(diff.getDroppedColumns()).isEmpty(); + assertThat(diff.hasTypeChanges()).isFalse(); + } + + @Test + void detectsRenameWithTypeChange() { + RowType oldRow = rowType(ID, NAME); + RowType newRow = rowType(ID, new RowType.RowField("full_name", new BigIntType(true))); + + SchemaDiff diff = SchemaDiff.compute(oldRow, newRow); + + assertThat(diff.getRenames()).hasSize(1); + SchemaDiff.Rename rename = diff.getRenames().get(0); + assertThat(rename.getOldName()).isEqualTo("name"); + assertThat(rename.getNewName()).isEqualTo("full_name"); + assertThat(rename.getNewType()).isInstanceOf(BigIntType.class); + assertThat(diff.hasTypeChanges()).isTrue(); + } + + @Test + void exposesTypeChangeMapping() { + RowType oldRow = rowType(ID, NAME); + RowType newRow = rowType(ID, new RowType.RowField("name", new BigIntType(true))); + + SchemaDiff diff = SchemaDiff.compute(oldRow, newRow); + + assertThat(diff.getTypeChangeByColumn()).containsKey("name"); + assertThat(diff.getTypeChangeByColumn().get("name")).isInstanceOf(BigIntType.class); + } + + @Test + void rejectsUnsafeDropAndAddAtDifferentPositions() { + // 'name' removed at position 1; 'nickname' added at position 2. Not position-aligned, + // so it cannot be a pure RENAME and must be rejected rather than treated as DROP+ADD. + RowType oldRow = rowType(ID, NAME, AGE); + RowType newRow = rowType(ID, AGE, + new RowType.RowField("nickname", new VarCharType(true, VarCharType.MAX_LENGTH))); + + assertThatThrownBy(() -> SchemaDiff.compute(oldRow, newRow)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("RENAME"); + } +}