From 835847380d28763de9a0517d381fee20e8691a69 Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 15 Sep 2026 10:24:43 +0800 Subject: [PATCH 01/11] review-fixes(upsert-sink): resolve architecture-review must-fix items Address the A/B-level items from the architecture review of the upsert-delete-alter-table PR: A1+A3+A5: rewrite LanceUpsertSink first-write path - Replace the createDataset (Overwrite) branch with an open-or-create strategy (Dataset.open, fallback Dataset.create) that is atomic in the Lance native layer, eliminating multi-subtask Overwrite clobbering. - Remove the Files.exists() checks that silently misclassified remote paths (s3://, tbdsfs://) as non-existent. - OVERWRITE mode now explicitly rejects remote storage in the sink and documents that truncation must happen at DDL time. A2: tighten checkpoint / consistency model - Apply deletes before upserts within a flush so a half-failed flush never leaves a stale row that should have been superseded. - Drop the implicit flush() from close(); persistence boundary is now strictly the checkpoint, matching at-least-once semantics. - Reject NULL primary-key values and non-finite float PK values in the DELETE predicate (previously they produced silently no-op predicates). - Escalate OVERWRITE local-directory delete failures from LOG.warn to IOException (was B6). B2: bound the in-memory buffer - invoke() now triggers an early flush() once the collapsed key count reaches write.batch-size, preventing unbounded heap growth between checkpoints. Happens between events, so the delete-before-upsert invariant is preserved. B5: protect foreign-namespaced dataset config during ALTER - applyTableProperties() no longer UNSETs config keys that look namespaced (contain a dot), preserving metadata written by sister engines (Spark, Trino, Ray) across a Flink ALTER TABLE ... RESET. Engineering-quality - LanceDynamicTableSink assigns a stable operator UID + name to the upsert sink so state mapping survives job upgrades / savepoints. - Extract PrimaryKeySelector.project() and have LanceUpsertSink. extractKey() delegate to it, so the keyBy routing key and the buffer key can never drift. - Reject comma-containing primary-key column names in PrimaryKeyPersistence.persist() to keep the comma-delimited encoding unambiguous. - Narrow ArrowArrayStreams from public to package-private. - LanceUpsertSinkITCase.twoSubtasksConcurrentFirstWrite() is now truly concurrent (CountDownLatch double-barrier + ExecutorService). - New LanceUpsertSinkITCase.closeDoesNotImplicitlyFlush() regression test for the A2 close() semantics. Not addressed in this commit (tracked in .gh-comments/): - A4 typed DELETE encoding (replace SQL string with mergeInsert WhenMatched.Delete or IN-list). - A6 consume TableChange in LanceCatalog.alterTable instead of SchemaDiff heuristic. - A7 single source of truth for connector option classification. - B1 hot-key metrics; B3 RootAllocator cap (needs cross-component coordination); B4 cross-engine PK metadata key. mvn -o clean test-compile: BUILD SUCCESS (0 errors) on 1.18/1.19/1.20. --- .../connector/lance/ArrowArrayStreams.java | 4 +- .../connector/lance/LanceUpsertSink.java | 298 +++++++++++------- .../lance/PrimaryKeyPersistence.java | 19 +- .../connector/lance/PrimaryKeySelector.java | 12 + .../connector/lance/table/LanceCatalog.java | 30 +- .../lance/table/LanceDynamicTableSink.java | 9 +- .../lance/LanceUpsertSinkITCase.java | 90 +++++- 7 files changed, 334 insertions(+), 128 deletions(-) diff --git a/src/main/java/org/apache/flink/connector/lance/ArrowArrayStreams.java b/src/main/java/org/apache/flink/connector/lance/ArrowArrayStreams.java index eb0d671..279c66e 100644 --- a/src/main/java/org/apache/flink/connector/lance/ArrowArrayStreams.java +++ b/src/main/java/org/apache/flink/connector/lance/ArrowArrayStreams.java @@ -43,7 +43,7 @@ * 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 { +final class ArrowArrayStreams { private ArrowArrayStreams() { // utility class @@ -58,7 +58,7 @@ private ArrowArrayStreams() { * @param root populated data batch * @return the merge-insert result */ - public static MergeInsertResult mergeInsert( + static MergeInsertResult mergeInsert( Dataset dataset, MergeInsertParams params, BufferAllocator allocator, diff --git a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java index cb67c46..287aa0a 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java @@ -26,7 +26,6 @@ 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; @@ -41,14 +40,9 @@ 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; @@ -57,11 +51,12 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.net.URI; +import java.net.URISyntaxException; 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; @@ -70,9 +65,9 @@ * 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 +I}/{@code +U}/{@code -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: * *

+ * + *

Consistency model

+ *

This sink provides at-least-once semantics, not exactly-once: + *

+ * + *

Concurrency and first-write

+ *

The sink uses an open-or-create strategy in {@link #open}: it first tries + * {@link Dataset#open}, and if that fails it falls back to + * {@link Dataset#create(BufferAllocator, String, Schema, WriteParams)} which is atomic in the + * Lance native layer. Concurrent first-writes from multiple subtasks therefore either observe + * the same dataset (both {@code open} succeeds) or race on {@code create} (one wins, others fall + * back to {@code open}); either way no {@code Overwrite} clobber can occur. + * + *

This deliberately replaces the previous {@code Files.exists()} check, which is only correct + * for the local filesystem and would silently mis-classify remote paths (s3://, tbdsfs://) as + * non-existent, causing every subtask to re-create and clobber the dataset. + * + *

Hot keys and buffering

+ *

Because {@code keyBy} routes all events for a given primary key to a single subtask, a + * skewed key distribution (one dominant key) will bottleneck the whole pipeline on that + * subtask. Two-level hashing is not applicable — it would break in-key ordering, which the + * delete-before-upsert per-flush invariant depends on. Callers with a known-skewed natural PK + * should salt or composite the key at the SQL layer. + * + *

The per-key buffer is bounded by {@code write.batch-size}: once the collapsed key count + * reaches the threshold, {@link #flush} is invoked between events (never mid-flush), keeping + * heap usage predictable in the gap between checkpoints. */ public class LanceUpsertSink extends RichSinkFunction implements CheckpointedFunction { @@ -91,6 +126,7 @@ public class LanceUpsertSink extends RichSinkFunction implements Checkp private final RowType rowType; private final List primaryKeys; private final int[] keyIndices; + private final LogicalType[] keyTypes; private transient BufferAllocator allocator; private transient Dataset dataset; @@ -104,13 +140,24 @@ public LanceUpsertSink(LanceOptions options, RowType rowType, List prima this.rowType = rowType; this.primaryKeys = primaryKeys; this.keyIndices = keyIndices; + // Pre-resolve key types once so extractKey doesn't dispatch through rowType per event, + // and so the buffer key and the keyBy routing key (PrimaryKeySelector) share the exact + // same projection logic (see PrimaryKeySelector#project). + this.keyTypes = new LogicalType[keyIndices.length]; + for (int i = 0; i < keyIndices.length; i++) { + this.keyTypes[i] = rowType.getTypeAt(keyIndices[i]); + } } @Override public void open(Configuration parameters) throws Exception { super.open(parameters); - LOG.info("Opening Lance Upsert Sink: {}", options.getPath()); + String datasetPath = options.getPath(); + if (datasetPath == null || datasetPath.isEmpty()) { + throw new IllegalArgumentException("Lance dataset path cannot be empty"); + } + LOG.info("Opening Lance Upsert Sink: {}", datasetPath); this.allocator = new RootAllocator(Long.MAX_VALUE); this.buffer = new LinkedHashMap<>(); @@ -118,48 +165,95 @@ public void open(Configuration parameters) throws Exception { 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"); + // OVERWRITE mode: only supported for local filesystem paths. For remote storage the + // truncation must be performed at DDL time (e.g. via catalog CREATE OR REPLACE) — the + // sink cannot reliably drop a remote dataset without an SDK-provided API. + if (options.getWriteMode() == LanceOptions.WriteMode.OVERWRITE) { + if (isLocalPath(datasetPath)) { + Path path = Paths.get(datasetPath); + if (Files.exists(path)) { + LOG.info("Overwrite mode, deleting existing local dataset: {}", datasetPath); + deleteDirectory(path); + } + } else { + throw new UnsupportedOperationException( + "write.mode=overwrite is not supported for remote storage in the upsert " + + "sink. Please drop/recreate the table via the catalog before writing. " + + "Path: " + datasetPath); + } } - Path path = Paths.get(datasetPath); - boolean datasetExists = Files.exists(path); + // Open-or-create: atomic in the Lance native layer, safe under concurrent first-writes. + this.dataset = openOrCreate(datasetPath); - 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); - } + // Persist the primary keys into the dataset config. Idempotent; safe to call on every open. + PrimaryKeyPersistence.persist(dataset, primaryKeys); LOG.info("Lance Upsert Sink opened, primary keys: {}", primaryKeys); } + /** + * Return an open {@link Dataset} handle, creating an empty dataset first if it does not yet + * exist. This unifies first-write and steady-state paths so that both go through + * {@code mergeInsert}/{@code delete}, eliminating the previous {@code Overwrite}-based + * first-write that could clobber peer subtasks' commits. + */ + private Dataset openOrCreate(String datasetPath) throws IOException { + try { + return Dataset.open(datasetPath, allocator); + } catch (Exception openFailure) { + LOG.debug("Dataset.open failed ({}), attempting create for: {}", + openFailure.getMessage(), datasetPath); + try { + return Dataset.create( + allocator, datasetPath, arrowSchema, new WriteParams.Builder().build()); + } catch (Exception createFailure) { + // A peer subtask likely won the create race; try open once more. + try { + return Dataset.open(datasetPath, allocator); + } catch (Exception reopenFailure) { + IOException io = new IOException( + "Failed to open or create Lance dataset: " + datasetPath, reopenFailure); + io.addSuppressed(createFailure); + io.addSuppressed(openFailure); + throw io; + } + } + } + } + @Override - public void invoke(RowData value, Context context) { + public void invoke(RowData value, Context context) throws IOException { 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; + return; default: LOG.warn("Ignoring unsupported RowKind: {}", kind); + return; + } + // Bound the buffer: without this, a large gap between checkpoints combined with a wide + // key space produces unbounded heap growth. This early flush is safe because it happens + // between events (never mid-flush), so the delete-before-upsert invariant per flush and + // the per-key collapsing invariant within one flush are both preserved. + if (buffer.size() >= options.getWriteBatchSize()) { + flush(); } } /** * Flush the collapsed per-key buffer to Lance. + * + *

Deletes are applied before upserts within a flush: this guarantees that if the + * flush half-fails, the target never contains a stale row that should have been superseded. + * Combined with per-key collapsing (only the last event per key is kept), the operation is + * idempotent under upstream replay. */ public void flush() throws IOException { if (buffer.isEmpty()) { @@ -176,69 +270,21 @@ public void flush() throws IOException { } } - 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(); + // Delete first, so a partial failure never leaves a row we intended to remove. + if (!deletes.isEmpty()) { + deleteRows(deletes); } - - 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; + if (!upserts.isEmpty()) { + mergeInsertRows(upserts); } - 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); - } + totalWrittenRows += upserts.size() + deletes.size(); + buffer.clear(); } /** - * Apply native upsert via {@code mergeInsert}. + * Apply native upsert via {@code mergeInsert}. Also handles the "insert into empty dataset" + * case, since the dataset was materialized empty in {@link #open}. */ private void mergeInsertRows(List rows) throws IOException { MergeInsertParams params = new MergeInsertParams(primaryKeys) @@ -273,6 +319,14 @@ private String buildDeletePredicate(List rows) { String column = rowType.getFieldNames().get(keyIndex); LogicalType type = rowType.getTypeAt(keyIndex); Object value = RowDataFieldAccessor.readField(row, keyIndex, type); + if (value == null) { + // NULL primary keys cannot participate in an equality predicate + // (col = NULL is UNKNOWN in SQL, so the row would never be matched + // and the delete would silently no-op). Reject explicitly. + throw new IllegalStateException( + "NULL primary-key value is not supported for DELETE on column '" + + column + "'"); + } ands.add(column + " = " + formatSqlValue(value, type)); } ors.add("(" + String.join(" AND ", ands) + ")"); @@ -284,15 +338,27 @@ private String buildDeletePredicate(List rows) { * 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 FloatType) { + float f = (Float) value; + if (!Float.isFinite(f)) { + throw new IllegalStateException( + "Non-finite float primary-key value (" + f + ") is not supported for DELETE"); + } + return Float.toString(f); + } + if (type instanceof DoubleType) { + double d = (Double) value; + if (!Double.isFinite(d)) { + throw new IllegalStateException( + "Non-finite double primary-key value (" + d + ") is not supported for DELETE"); + } + return Double.toString(d); + } if (type instanceof VarCharType) { StringData stringData = (StringData) value; return "'" + stringData.toString().replace("'", "''") + "'"; @@ -304,18 +370,19 @@ private String formatSqlValue(Object value, LogicalType type) { /** * Project the primary-key columns into a key that honors equals/hashCode. */ + /** + * Project the primary-key columns into a key that honors equals/hashCode. Delegates to + * {@link PrimaryKeySelector#project} so the buffer key here is byte-for-byte identical to + * the {@code keyBy} routing key upstream. + */ 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; + return PrimaryKeySelector.project(value, keyIndices, keyTypes); } @Override public void snapshotState(FunctionSnapshotContext context) throws Exception { LOG.debug("Snapshot state, checkpointId: {}", context.getCheckpointId()); + // Persistence boundary: only checkpoint triggers a flush. flush(); } @@ -327,11 +394,9 @@ public void initializeState(FunctionInitializationContext context) { @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); - } + // Do NOT flush here: close() runs on cancel/restart as well; writing on those paths + // would violate the "checkpoint is the persistence boundary" contract. Rows still in + // buffer are dropped and will be re-delivered by the (replayable) source on restart. if (dataset != null) { try { dataset.close(); @@ -352,15 +417,32 @@ public void close() throws Exception { super.close(); } + /** + * Whether the given path denotes the local filesystem (as opposed to s3://, tbdsfs://, hdfs:// + * etc.). A path with no scheme, or the explicit {@code file:} scheme, is considered local. + */ + private static boolean isLocalPath(String path) { + try { + URI uri = new URI(path); + String scheme = uri.getScheme(); + return scheme == null || "file".equalsIgnoreCase(scheme); + } catch (URISyntaxException e) { + // A parse failure means it's almost certainly a plain local path. + return true; + } + } + 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); - } - }); + try (java.util.stream.Stream children = Files.list(path)) { + children.forEach(child -> { + try { + deleteDirectory(child); + } catch (IOException e) { + throw new RuntimeException("Failed to delete file: " + child, e); + } + }); + } } Files.deleteIfExists(path); } diff --git a/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java b/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java index 0311da1..c862e16 100644 --- a/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java +++ b/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java @@ -21,7 +21,6 @@ import org.lance.Dataset; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Map; @@ -33,6 +32,11 @@ *

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()}. + * + *

Encoding limitation: column names are joined with a literal comma. Names containing + * a comma are rejected in {@link #persist} to keep round-trip decoding unambiguous. Flink / + * Arrow / Lance schema tooling in practice constrain identifiers to a comma-free character + * class, so this restriction is not observable in normal usage. */ public final class PrimaryKeyPersistence { @@ -48,11 +52,24 @@ private PrimaryKeyPersistence() { * * @param dataset the Lance dataset (must already be materialized) * @param primaryKeys ordered primary-key column names; empty/null writes nothing + * @throws IllegalArgumentException if any column name contains a comma (which would make + * the comma-delimited encoding ambiguous on load) */ public static void persist(Dataset dataset, List primaryKeys) { if (dataset == null || primaryKeys == null || primaryKeys.isEmpty()) { return; } + for (String column : primaryKeys) { + if (column == null || column.isEmpty()) { + throw new IllegalArgumentException( + "Primary-key column names must be non-empty; got: " + primaryKeys); + } + if (column.indexOf(',') >= 0) { + throw new IllegalArgumentException( + "Primary-key column name must not contain a comma (would corrupt the " + + "comma-delimited encoding): '" + column + "'"); + } + } dataset.updateConfig(Collections.singletonMap(PK_CONFIG_KEY, String.join(",", primaryKeys))); } diff --git a/src/main/java/org/apache/flink/connector/lance/PrimaryKeySelector.java b/src/main/java/org/apache/flink/connector/lance/PrimaryKeySelector.java index 4b5a97d..14cbc83 100644 --- a/src/main/java/org/apache/flink/connector/lance/PrimaryKeySelector.java +++ b/src/main/java/org/apache/flink/connector/lance/PrimaryKeySelector.java @@ -50,6 +50,18 @@ public PrimaryKeySelector(int[] keyIndices, LogicalType[] keyTypes) { @Override public RowData getKey(RowData value) { + return project(value, keyIndices, keyTypes); + } + + /** + * Project the primary-key columns out of a {@link RowData} into a freshly allocated + * {@link GenericRowData} that honors {@code equals}/{@code hashCode}. + * + *

Centralizing the projection here keeps the {@code keyBy} routing key (this class) and + * the sink's in-memory buffer key ({@code LanceUpsertSink}) byte-for-byte identical; any + * drift between the two would silently break the "same PK routed to same subtask" contract. + */ + public static RowData project(RowData value, int[] keyIndices, LogicalType[] keyTypes) { GenericRowData key = new GenericRowData(keyIndices.length); for (int i = 0; i < keyIndices.length; i++) { key.setField(i, RowDataFieldAccessor.readField(value, keyIndices[i], keyTypes[i])); 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 c4630c4..71f03bf 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 @@ -735,20 +735,46 @@ private void applyTableProperties(Dataset dataset, CatalogBaseTable newTable) { dataset.updateConfig(toSet); } + // UNSET is applied conservatively: we only remove keys that Flink itself might have + // written. Any config key using a namespaced form ("engine.category.name") is treated + // as potentially owned by a sister engine (Spark / Trino / Ray) and left untouched, + // even if it looks like a user TBLPROPERTY by our own classifier. Keys the user + // explicitly re-sets in this ALTER stay writable via the toSet path above. + // + // TODO(follow-up A7): move to a "Flink writes under flink.* namespace only" model so + // this heuristic can be replaced with an exact prefix match. 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 (!isTblProperty(key)) { + continue; + } + if (options.containsKey(key)) { + continue; + } + if (isForeignNamespacedKey(key)) { + LOG.debug("Skipping UNSET of foreign-namespaced config key: {}", key); + continue; } + toUnset.add(key); } if (!toUnset.isEmpty()) { dataset.deleteConfigKeys(toUnset); } } + /** + * Whether the given config key looks like it was written by another engine (Spark, Trino, + * Ray, ...). We use a simple heuristic: a key containing a dot ({@code engine.some.prop}) + * is treated as namespaced and therefore not ours to delete. The Flink primary-key metadata + * key ({@link PrimaryKeyPersistence#PK_CONFIG_KEY}) is filtered separately by the caller. + */ + private static boolean isForeignNamespacedKey(String key) { + return key != null && key.indexOf('.') >= 0; + } + private boolean isTblProperty(String key) { if (key == null) { return false; 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 3490200..c012358 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 @@ -108,11 +108,14 @@ public SinkRuntimeProvider getSinkRuntimeProvider(Context context) { @Override public DataStreamSink consumeDataStream( ProviderContext providerContext, DataStream dataStream) { - DataStream keyed = - dataStream.keyBy(new PrimaryKeySelector(primaryKeyIndices, primaryKeyTypes)); + DataStream keyed = dataStream + .keyBy(new PrimaryKeySelector(primaryKeyIndices, primaryKeyTypes)); LanceUpsertSink upsertSink = new LanceUpsertSink(options, rowType, primaryKeys, primaryKeyIndices); - return keyed.addSink(upsertSink); + DataStreamSink sink = keyed.addSink(upsertSink).name("LanceUpsertSink"); + // Assign a stable UID so state mapping survives job upgrades / savepoints. + providerContext.generateUid("lance-upsert-sink").ifPresent(sink::uid); + return sink; } }; } diff --git a/src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkITCase.java b/src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkITCase.java index e282040..454e89d 100644 --- a/src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkITCase.java +++ b/src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkITCase.java @@ -39,6 +39,12 @@ import java.nio.file.Path; import java.util.Arrays; import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import static org.assertj.core.api.Assertions.assertThat; @@ -290,25 +296,85 @@ void twoSubtasksConcurrentWriteToExistingDataset() throws Exception { 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(), + // Dataset does NOT exist yet. Both subtasks open + flush in parallel: the previous + // Overwrite-based first-write would race here and clobber one subtask's rows. With the + // open-or-create + mergeInsert path, both writes must land. + final LanceUpsertSink s1 = new LanceUpsertSink(options(path), rowType(), Arrays.asList(PRIMARY_KEYS), KEY_INDICES); - LanceUpsertSink s2 = new LanceUpsertSink(options(path), rowType(), + final LanceUpsertSink s2 = new LanceUpsertSink(options(path), rowType(), Arrays.asList(PRIMARY_KEYS), KEY_INDICES); - s1.open(new Configuration()); - s2.open(new Configuration()); + final CountDownLatch openBarrier = new CountDownLatch(2); + final CountDownLatch flushBarrier = new CountDownLatch(2); + final AtomicReference failure = new AtomicReference<>(); + + ExecutorService pool = Executors.newFixedThreadPool(2); try { - s1.invoke(row(1L, "one", RowKind.INSERT), null); - s1.flush(); - s2.invoke(row(2L, "two", RowKind.INSERT), null); - s2.flush(); + Future f1 = pool.submit(() -> runSubtask( + s1, 1L, "one", openBarrier, flushBarrier, failure)); + Future f2 = pool.submit(() -> runSubtask( + s2, 2L, "two", openBarrier, flushBarrier, failure)); + f1.get(30, TimeUnit.SECONDS); + f2.get(30, TimeUnit.SECONDS); } finally { - s1.close(); - s2.close(); + pool.shutdownNow(); + try { s1.close(); } catch (Exception ignored) { /* best effort */ } + try { s2.close(); } catch (Exception ignored) { /* best effort */ } + } + + if (failure.get() != null) { + throw new AssertionError("Concurrent first-write failed", failure.get()); } - // The second first-write must not clobber the first (regression: Overwrite clobbering). + // The concurrent first-write must not lose either row (regression: Overwrite clobbering). assertThat(countRows(path)).isEqualTo(2L); } + + private void runSubtask( + LanceUpsertSink sink, + long id, + String name, + CountDownLatch openBarrier, + CountDownLatch flushBarrier, + AtomicReference failure) { + try { + sink.open(new Configuration()); + // Rendezvous after open so both subtasks contend on the empty dataset. + openBarrier.countDown(); + openBarrier.await(30, TimeUnit.SECONDS); + + sink.invoke(row(id, name, RowKind.INSERT), null); + // Rendezvous again so the two flushes hit mergeInsert concurrently. + flushBarrier.countDown(); + flushBarrier.await(30, TimeUnit.SECONDS); + + sink.flush(); + } catch (Throwable t) { + failure.compareAndSet(null, t); + // Release peers so the test doesn't hang on the barrier. + openBarrier.countDown(); + flushBarrier.countDown(); + } + } + + @Test + @DisplayName("close() must not implicitly flush uncheckpointed rows") + void closeDoesNotImplicitlyFlush() throws Exception { + String path = tempDir.resolve("no_close_flush").toString(); + LanceUpsertSink sink = new LanceUpsertSink(options(path), rowType(), + Arrays.asList(PRIMARY_KEYS), KEY_INDICES); + + sink.open(new Configuration()); + try { + // Enqueue a row but do NOT flush. close() must drop it (checkpoint is the + // persistence boundary; the source is expected to replay on restart). + sink.invoke(row(42L, "unflushed", RowKind.INSERT), null); + } finally { + sink.close(); + } + + // Dataset was materialized empty in open() and never received a checkpoint-driven + // flush, so it must contain zero rows. + assertThat(countRows(path)).isZero(); + } } From fa43be799f08b36774b2a740e05cef9bea17ea51 Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 11:05:05 +0800 Subject: [PATCH 02/11] fix: address PR #76 review follow-ups A4, A6, A7, B1, B3 A4: rewrite DELETE on key-only mergeInsert Replace the SQL-string delete predicate with a key-only mergeInsert using withMatchedDelete + WhenNotMatched.DoNothing. This closes the type-coverage gap (DATE/TIME/TIMESTAMP/DECIMAL/VARBINARY), removes the injection surface from unescaped column names, and eliminates the O(N) predicate growth. Also adopt the post-merge dataset handle: mergeInsert commits a new version and returns a new handle rather than mutating the receiver, so the sink was reading a stale snapshot after every merge. A6: drive ALTER TABLE from TableChange Override the alterTable overload that receives the planner's explicit change list instead of inferring intent from a schema diff. A DROP followed by an ADD of the same name is now applied literally; the diff heuristic could not distinguish it from a rename and refused the statement. Changes are applied rename -> drop -> add, since re-creating a column whose name still exists is rejected by Lance as a type conflict. ALTER COLUMN type changes stay rejected because the SDK's castTo is a verified silent no-op. The diff-based path remains as a fallback when no changes are supplied. A7: derive option classification from the factories Introduce LanceOptionRegistry, which builds the reserved-key set from the factories' live ConfigOption declarations, and delete the hand-maintained copy in LanceCatalog. The two had already drifted: s3-virtual-hosted-style and s3-allow-http were declared by LanceCatalogFactory but missing from the copy, so they leaked into the dataset config as user properties. Scope RESET to Flink-owned keys (the flink.* namespace plus unnamespaced keys) instead of treating every dotted key as foreign. The old rule was safe against cross-engine deletion but also made RESET impossible for dotted properties Flink itself wrote. B1: document the hot-primary-key constraint keyBy on the primary key is required for flush ordering, so per-key throughput is bounded by one subtask. Document composite or salted keys as the mitigation and record why two-level hashing is not offered. B3: make the Arrow allocator bound configurable Route all ten allocation sites through LanceAllocators and add the opt-in arrow.allocator-max-bytes option; unset preserves the previous unbounded behaviour. Each allocator is named so an OOM identifies its owner. --- docs/src/operations/dml/insert-into.md | 88 +++- .../connector/lance/LanceAggregateSource.java | 5 +- .../connector/lance/LanceIndexBuilder.java | 4 +- .../connector/lance/LanceInputFormat.java | 8 +- .../flink/connector/lance/LanceSink.java | 5 +- .../flink/connector/lance/LanceSource.java | 5 +- .../connector/lance/LanceUpsertSink.java | 191 +++++--- .../connector/lance/LanceVectorSearch.java | 4 +- .../connector/lance/RowDataFieldAccessor.java | 97 +++- .../connector/lance/config/LanceOptions.java | 39 ++ .../connector/lance/table/LanceCatalog.java | 213 +++++++-- .../table/LanceNamespaceCatalogFactory.java | 4 +- .../lance/table/LanceOptionRegistry.java | 165 +++++++ .../connector/lance/util/LanceAllocators.java | 129 ++++++ .../lance/ArrowArrayStreamsTestAccess.java | 51 ++ .../LanceUpsertSinkDeleteRewriteTest.java | 438 ++++++++++++++++++ .../lance/MergeInsertHandleStalenessTest.java | 213 +++++++++ .../table/LanceCatalogTableChangeITCase.java | 329 +++++++++++++ .../lance/table/LanceOptionRegistryTest.java | 241 ++++++++++ 19 files changed, 2103 insertions(+), 126 deletions(-) create mode 100644 src/main/java/org/apache/flink/connector/lance/table/LanceOptionRegistry.java create mode 100644 src/main/java/org/apache/flink/connector/lance/util/LanceAllocators.java create mode 100644 src/test/java/org/apache/flink/connector/lance/ArrowArrayStreamsTestAccess.java create mode 100644 src/test/java/org/apache/flink/connector/lance/LanceUpsertSinkDeleteRewriteTest.java create mode 100644 src/test/java/org/apache/flink/connector/lance/MergeInsertHandleStalenessTest.java create mode 100644 src/test/java/org/apache/flink/connector/lance/table/LanceCatalogTableChangeITCase.java create mode 100644 src/test/java/org/apache/flink/connector/lance/table/LanceOptionRegistryTest.java diff --git a/docs/src/operations/dml/insert-into.md b/docs/src/operations/dml/insert-into.md index 861a656..2818aee 100644 --- a/docs/src/operations/dml/insert-into.md +++ b/docs/src/operations/dml/insert-into.md @@ -1,7 +1,7 @@ # INSERT INTO -The Lance Flink sink appends rows to a Lance dataset. Write mode is controlled by -the `write.mode` option. +The Lance Flink sink writes rows to a Lance dataset. Behaviour depends on whether +the table declares a primary key. ## Write modes @@ -17,6 +17,32 @@ INSERT INTO vectors VALUES (1, 'Hello World', ARRAY[0.1, 0.2, 0.3, 0.4]); ``` +## Upsert mode + +When a table declares `PRIMARY KEY (...) NOT ENFORCED`, the sink accepts a CDC +changelog stream and maps it onto Lance native operations: + +| Change | Applied as | +|---|---| +| `+I` / `+U` | `mergeInsert` with update-all + insert-all | +| `-D` | key-only `mergeInsert` with matched-delete + not-matched-do-nothing | +| `-U` | dropped (the new value carries all required state) | + +```sql +CREATE TABLE users ( + id BIGINT, + name STRING, + PRIMARY KEY (id) NOT ENFORCED +) WITH ( + 'connector' = 'lance', + 'path' = '/data/users.lance' +); +``` + +Deletes are matched by primary-key value, so any column type the connector can +write can also serve as a primary key, including `DATE`, `TIMESTAMP`, `DECIMAL` +and `VARBINARY`. + ## Sink options | Option | Default | Description | @@ -24,6 +50,53 @@ INSERT INTO vectors VALUES | `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 | +| `arrow.allocator-max-bytes` | unlimited | Upper bound in bytes for the Arrow allocator | + +### Bounding Arrow memory + +Each sink, source, catalog and index builder creates its own Arrow allocator. +By default these are unbounded, which lets a single oversized batch exhaust +off-heap memory and affect other slots in the same TaskManager. + +`arrow.allocator-max-bytes` caps each allocator individually: + +```sql +CREATE TABLE users ( + id BIGINT, + name STRING, + PRIMARY KEY (id) NOT ENFORCED +) WITH ( + 'connector' = 'lance', + 'path' = '/data/users.lance', + 'arrow.allocator-max-bytes' = '536870912' +); +``` + +The bound is per allocator instance, not a total for the job. Size it against +one component's working set — roughly `write.batch-size` times the row width, +with headroom — rather than against the TaskManager's whole off-heap budget. + +## Hot primary keys + +In upsert mode events are routed with `keyBy` on the primary key so that all +events for a key reach one subtask in order. This ordering is required for +correctness, but it also means a single key's throughput is capped by one +subtask. + +If the key distribution is skewed, that subtask gates the whole pipeline while +its peers idle. The sink collapses repeated writes to the same key within a +checkpoint, which absorbs update-heavy skew, but not sheer volume concentrated +on one key. + +Where the natural key is known to be skewed, prefer a composite primary key +including a higher-cardinality column, or salt the key: + +```sql +PRIMARY KEY (tenant_id, event_id) NOT ENFORCED +``` + +Two-level hashing is not offered: it would break in-key ordering, which the +sink's flush sequencing depends on. ## Current limitations @@ -31,9 +104,10 @@ INSERT INTO vectors VALUES |---|---| | `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 | +| Primary key / upsert | ✅ | +| `DELETE` (via CDC changelog) | ✅ | +| `UPDATE` (standalone statement) | ❌ — not implemented | -> The sink currently declares insert-only changelog mode. CDC `UPDATE` / `DELETE` -> support is tracked in the connector roadmap. +> Standalone `UPDATE` and `DELETE` SQL statements are not supported; row-level +> changes are applied through a CDC changelog stream into a table declaring a +> primary key. diff --git a/src/main/java/org/apache/flink/connector/lance/LanceAggregateSource.java b/src/main/java/org/apache/flink/connector/lance/LanceAggregateSource.java index d2054a8..15d2aa2 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceAggregateSource.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceAggregateSource.java @@ -34,7 +34,7 @@ import org.lance.ipc.LanceScanner; import org.lance.ipc.ScanOptions; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; +import org.apache.flink.connector.lance.util.LanceAllocators; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; @@ -103,7 +103,8 @@ public void open(Configuration parameters) throws Exception { LOG.info("Aggregate info: {}", aggregateInfo); this.running = true; - this.allocator = new RootAllocator(Long.MAX_VALUE); + this.allocator = LanceAllocators.create( + "lance-aggregate-source", options.getArrowAllocatorMaxBytes()); // Open Lance dataset String datasetPath = options.getPath(); diff --git a/src/main/java/org/apache/flink/connector/lance/LanceIndexBuilder.java b/src/main/java/org/apache/flink/connector/lance/LanceIndexBuilder.java index 2c5639b..793c909 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceIndexBuilder.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceIndexBuilder.java @@ -29,7 +29,7 @@ import org.lance.index.vector.PQBuildParams; import org.lance.index.vector.VectorIndexParams; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; +import org.apache.flink.connector.lance.util.LanceAllocators; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -104,7 +104,7 @@ public IndexBuildResult buildIndex() throws IOException { try { // Initialize resources - this.allocator = new RootAllocator(Long.MAX_VALUE); + this.allocator = LanceAllocators.create("lance-index-builder"); this.dataset = Dataset.open(datasetPath, allocator); // Validate column exists diff --git a/src/main/java/org/apache/flink/connector/lance/LanceInputFormat.java b/src/main/java/org/apache/flink/connector/lance/LanceInputFormat.java index 8177138..9972314 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceInputFormat.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceInputFormat.java @@ -34,7 +34,7 @@ import org.lance.ipc.LanceScanner; import org.lance.ipc.ScanOptions; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; +import org.apache.flink.connector.lance.util.LanceAllocators; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; @@ -106,7 +106,8 @@ public LanceSplit[] createInputSplits(int minNumSplits) throws IOException { throw new IOException("Dataset path cannot be empty"); } - BufferAllocator tempAllocator = new RootAllocator(Long.MAX_VALUE); + BufferAllocator tempAllocator = LanceAllocators.create( + "lance-input-format-probe", options.getArrowAllocatorMaxBytes()); try { // Honor read.version / read.as-of-timestamp for time-travel reads (issue #5). Dataset tempDataset = LanceOpener.open(datasetPath, tempAllocator, options); @@ -139,7 +140,8 @@ public InputSplitAssigner getInputSplitAssigner(LanceSplit[] inputSplits) { public void open(LanceSplit split) throws IOException { LOG.info("Opening split: {}", split); - this.allocator = new RootAllocator(Long.MAX_VALUE); + this.allocator = LanceAllocators.create( + "lance-input-format", options.getArrowAllocatorMaxBytes()); this.reachedEnd = false; // Open dataset 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 ef3757e..b3635f0 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceSink.java @@ -38,7 +38,7 @@ import org.lance.operation.Append; import org.lance.operation.Overwrite; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; +import org.apache.flink.connector.lance.util.LanceAllocators; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; @@ -103,7 +103,8 @@ public void open(Configuration parameters) throws Exception { LOG.info("Opening Lance Sink: {}", options.getPath()); - this.allocator = new RootAllocator(Long.MAX_VALUE); + this.allocator = LanceAllocators.create( + "lance-sink", options.getArrowAllocatorMaxBytes()); this.buffer = new ArrayList<>(options.getWriteBatchSize()); this.totalWrittenRows = 0; this.isFirstWrite = true; diff --git a/src/main/java/org/apache/flink/connector/lance/LanceSource.java b/src/main/java/org/apache/flink/connector/lance/LanceSource.java index 12d1614..568e633 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceSource.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceSource.java @@ -33,7 +33,7 @@ import org.lance.ipc.LanceScanner; import org.lance.ipc.ScanOptions; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; +import org.apache.flink.connector.lance.util.LanceAllocators; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; @@ -117,7 +117,8 @@ public void open(Configuration parameters) throws Exception { this.running = true; this.emittedCount = 0; - this.allocator = new RootAllocator(Long.MAX_VALUE); + this.allocator = LanceAllocators.create( + "lance-source", options.getArrowAllocatorMaxBytes()); // Open Lance dataset String datasetPath = options.getPath(); diff --git a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java index 287aa0a..fd1bcd2 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java @@ -27,25 +27,18 @@ import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction; import org.apache.flink.streaming.api.functions.sink.RichSinkFunction; 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.Dataset; import org.lance.WriteParams; import org.lance.merge.MergeInsertParams; +import org.lance.merge.MergeInsertResult; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; +import org.apache.flink.connector.lance.util.LanceAllocators; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -73,10 +66,28 @@ *

  • {@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 DELETE} → native key-only {@link Dataset#mergeInsert} + * ({@code WhenMatched.Delete} + {@code WhenNotMatched.DoNothing}); see + * {@link #deleteRows}.
    • *
  • {@code UPDATE_BEFORE} → dropped (upsert has no need for the old value).
    • * * + *

    Hot primary keys

    + *

    {@code LanceDynamicTableSink} routes CDC events through {@code keyBy(PrimaryKeySelector)} so + * every event for a key lands on one subtask in order. That ordering is what makes the per-flush + * delete/upsert sequence correct, but it also means throughput for a key is bounded by a single + * subtask. If the key distribution is skewed — one {@code tenant_id} carrying most of the traffic, + * say — that subtask becomes the pipeline's ceiling while its peers idle. + * + *

    The per-key buffer collapsing below partially absorbs this for update-heavy workloads, since + * repeated writes to one key within a checkpoint fold into a single operation. It does not help + * when the hot key receives many distinct keys' worth of volume. + * + *

    Mitigation is schema-level: prefer a composite primary key that includes a + * higher-cardinality column, or salt the key, when the natural key is known to be skewed. + * Two-level hashing is deliberately not offered — it would break in-key ordering and with it the + * correctness of the flush sequence. + * *

    Consistency model

    *

    This sink provides at-least-once semantics, not exactly-once: *

    @@ -75,6 +81,14 @@ public class LanceTypeConverter implements Serializable { private static final long serialVersionUID = 1L; private static final Logger LOG = LoggerFactory.getLogger(LanceTypeConverter.class); + /** + * Bit width used for Arrow decimals. + * + *

    Flink's DECIMAL tops out at precision 38, which fits Decimal128, so widening to 256 is + * never required. + */ + private static final int DECIMAL_BIT_WIDTH = 128; + /** * Convert Arrow Schema to Flink RowType * @@ -154,11 +168,22 @@ public static LogicalType arrowTypeToFlinkType(Field field) { return new BinaryType(nullable, fixedBinary.getByteWidth()); } else if (arrowType instanceof ArrowType.Date) { return new DateType(nullable); + } else if (arrowType instanceof ArrowType.Time) { + ArrowType.Time timeType = (ArrowType.Time) arrowType; + return new TimeType(nullable, getTimestampPrecision(timeType.getUnit())); } else if (arrowType instanceof ArrowType.Timestamp) { ArrowType.Timestamp tsType = (ArrowType.Timestamp) arrowType; // Determine precision based on time unit int precision = getTimestampPrecision(tsType.getUnit()); + // A timezone marks an absolute instant, which is TIMESTAMP_LTZ on the Flink side. + // Without this split the zoned and unzoned forms would collapse into one. + if (tsType.getTimezone() != null) { + return new LocalZonedTimestampType(nullable, precision); + } return new TimestampType(nullable, precision); + } else if (arrowType instanceof ArrowType.Decimal) { + ArrowType.Decimal decimalType = (ArrowType.Decimal) arrowType; + return new DecimalType(nullable, decimalType.getPrecision(), decimalType.getScale()); } else if (arrowType instanceof ArrowType.FixedSizeList) { // Vector type: FixedSizeList ArrowType.FixedSizeList listType = (ArrowType.FixedSizeList) arrowType; @@ -228,10 +253,27 @@ public static Field flinkTypeToArrowField(String name, LogicalType logicalType) arrowType = new ArrowType.FixedSizeBinary(binaryType.getLength()); } else if (logicalType instanceof DateType) { arrowType = new ArrowType.Date(DateUnit.DAY); + } else if (logicalType instanceof TimeType) { + // Flink TIME is time-of-day without date. Arrow splits this across bit widths: + // Time32 carries SECOND/MILLISECOND, Time64 carries MICROSECOND/NANOSECOND. + TimeType timeType = (TimeType) logicalType; + TimeUnit timeUnit = getArrowTimeUnit(timeType.getPrecision()); + arrowType = new ArrowType.Time(timeUnit, getTimeBitWidth(timeUnit)); } else if (logicalType instanceof TimestampType) { TimestampType tsType = (TimestampType) logicalType; TimeUnit timeUnit = getArrowTimeUnit(tsType.getPrecision()); arrowType = new ArrowType.Timestamp(timeUnit, null); + } else if (logicalType instanceof LocalZonedTimestampType) { + // TIMESTAMP_LTZ denotes an absolute instant. Tagging the Arrow type with UTC keeps it + // distinguishable from a plain TIMESTAMP, which is what makes the round-trip lossless. + LocalZonedTimestampType ltzType = (LocalZonedTimestampType) logicalType; + TimeUnit timeUnit = getArrowTimeUnit(ltzType.getPrecision()); + arrowType = new ArrowType.Timestamp(timeUnit, "UTC"); + } else if (logicalType instanceof DecimalType) { + DecimalType decimalType = (DecimalType) logicalType; + arrowType = + new ArrowType.Decimal( + decimalType.getPrecision(), decimalType.getScale(), DECIMAL_BIT_WIDTH); } else if (logicalType instanceof ArrayType) { ArrayType arrayType = (ArrayType) logicalType; LogicalType elementType = arrayType.getElementType(); @@ -374,9 +416,17 @@ public static DataType toDataType(LogicalType logicalType) { return DataTypes.BINARY(binaryType.getLength()); } else if (logicalType instanceof DateType) { return DataTypes.DATE(); + } else if (logicalType instanceof TimeType) { + return DataTypes.TIME(((TimeType) logicalType).getPrecision()); } else if (logicalType instanceof TimestampType) { TimestampType tsType = (TimestampType) logicalType; return DataTypes.TIMESTAMP(tsType.getPrecision()); + } else if (logicalType instanceof LocalZonedTimestampType) { + LocalZonedTimestampType ltzType = (LocalZonedTimestampType) logicalType; + return DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(ltzType.getPrecision()); + } else if (logicalType instanceof DecimalType) { + DecimalType decimalType = (DecimalType) logicalType; + return DataTypes.DECIMAL(decimalType.getPrecision(), decimalType.getScale()); } else if (logicalType instanceof ArrayType) { ArrayType arrayType = (ArrayType) logicalType; DataType elementDataType = toDataType(arrayType.getElementType()); @@ -425,6 +475,22 @@ private static TimeUnit getArrowTimeUnit(int precision) { } } + /** + * Get the Arrow Time bit width required by a time unit. + * + *

    Arrow only allows Time32 for SECOND/MILLISECOND and Time64 for MICROSECOND/NANOSECOND; + * pairing a unit with the wrong width is rejected when the field is constructed. + */ + private static int getTimeBitWidth(TimeUnit timeUnit) { + switch (timeUnit) { + case SECOND: + case MILLISECOND: + return 32; + default: + return 64; + } + } + /** * Unsupported type exception */ diff --git a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java index 705bcee..992fcfd 100644 --- a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java +++ b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java @@ -19,6 +19,7 @@ package org.apache.flink.connector.lance.converter; import org.apache.flink.table.data.ArrayData; +import org.apache.flink.table.data.DecimalData; import org.apache.flink.table.data.GenericArrayData; import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; @@ -29,12 +30,15 @@ import org.apache.flink.table.types.logical.BinaryType; import org.apache.flink.table.types.logical.BooleanType; import org.apache.flink.table.types.logical.DateType; +import org.apache.flink.table.types.logical.DecimalType; 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.LocalZonedTimestampType; 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.TimeType; import org.apache.flink.table.types.logical.TimestampType; import org.apache.flink.table.types.logical.TinyIntType; import org.apache.flink.table.types.logical.VarBinaryType; @@ -44,15 +48,24 @@ import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TimeMicroVector; +import org.apache.arrow.vector.TimeMilliVector; +import org.apache.arrow.vector.TimeNanoVector; +import org.apache.arrow.vector.TimeSecVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TimeStampMilliTZVector; import org.apache.arrow.vector.TimeStampMilliVector; +import org.apache.arrow.vector.TimeStampNanoTZVector; import org.apache.arrow.vector.TimeStampNanoVector; +import org.apache.arrow.vector.TimeStampSecTZVector; import org.apache.arrow.vector.TimeStampSecVector; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.VarBinaryVector; @@ -66,6 +79,7 @@ import org.slf4j.LoggerFactory; import java.io.Serializable; +import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.time.Instant; import java.time.LocalDate; @@ -197,8 +211,17 @@ private Object readValue(FieldVector vector, int index, LogicalType logicalType) } else if (logicalType instanceof DateType) { int daysSinceEpoch = ((DateDayVector) vector).get(index); return daysSinceEpoch; + } else if (logicalType instanceof TimeType) { + return readTime(vector, index); } else if (logicalType instanceof TimestampType) { return readTimestamp(vector, index, (TimestampType) logicalType); + } else if (logicalType instanceof LocalZonedTimestampType) { + return readLocalZonedTimestamp(vector, index); + } else if (logicalType instanceof DecimalType) { + DecimalType decimalType = (DecimalType) logicalType; + BigDecimal value = ((DecimalVector) vector).getObject(index); + return DecimalData.fromBigDecimal( + value, decimalType.getPrecision(), decimalType.getScale()); } else if (logicalType instanceof ArrayType) { return readArray(vector, index, (ArrayType) logicalType); } else if (logicalType instanceof RowType) { @@ -234,6 +257,53 @@ private TimestampData readTimestamp(FieldVector vector, int index, TimestampType "Unsupported timestamp Vector type: " + vector.getClass().getSimpleName()); } + /** + * Read a TIME value as milliseconds since midnight. + * + *

    Flink represents TIME as an int holding milliseconds of the day, so the sub-millisecond + * units Arrow allows are narrowed down to that resolution here. + */ + private int readTime(FieldVector vector, int index) { + if (vector instanceof TimeSecVector) { + return ((TimeSecVector) vector).get(index) * 1000; + } else if (vector instanceof TimeMilliVector) { + return ((TimeMilliVector) vector).get(index); + } else if (vector instanceof TimeMicroVector) { + return (int) (((TimeMicroVector) vector).get(index) / 1000L); + } else if (vector instanceof TimeNanoVector) { + return (int) (((TimeNanoVector) vector).get(index) / 1_000_000L); + } + + throw new LanceTypeConverter.UnsupportedTypeException( + "Unsupported time Vector type: " + vector.getClass().getSimpleName()); + } + + /** + * Read a TIMESTAMP_LTZ value. + * + *

    Arrow exposes zoned timestamps through dedicated *TZ vectors, so these are distinct + * classes from the ones {@link #readTimestamp} handles. Values are epoch-based, which is + * exactly what {@link TimestampData#fromEpochMillis} expects, so no zone shifting is applied. + */ + private TimestampData readLocalZonedTimestamp(FieldVector vector, int index) { + if (vector instanceof TimeStampSecTZVector) { + return TimestampData.fromEpochMillis(((TimeStampSecTZVector) vector).get(index) * 1000L); + } else if (vector instanceof TimeStampMilliTZVector) { + return TimestampData.fromEpochMillis(((TimeStampMilliTZVector) vector).get(index)); + } else if (vector instanceof TimeStampMicroTZVector) { + long micros = ((TimeStampMicroTZVector) vector).get(index); + return TimestampData.fromEpochMillis( + Math.floorDiv(micros, 1000L), (int) Math.floorMod(micros, 1000L) * 1000); + } else if (vector instanceof TimeStampNanoTZVector) { + long nanos = ((TimeStampNanoTZVector) vector).get(index); + return TimestampData.fromEpochMillis( + Math.floorDiv(nanos, 1_000_000L), (int) Math.floorMod(nanos, 1_000_000L)); + } + + throw new LanceTypeConverter.UnsupportedTypeException( + "Unsupported zoned timestamp Vector type: " + vector.getClass().getSimpleName()); + } + /** * Read array value */ @@ -376,9 +446,17 @@ private Object getFieldValue(RowData rowData, int index, LogicalType logicalType return rowData.getBinary(index); } else if (logicalType instanceof DateType) { return rowData.getInt(index); + } else if (logicalType instanceof TimeType) { + return rowData.getInt(index); } else if (logicalType instanceof TimestampType) { TimestampType tsType = (TimestampType) logicalType; return rowData.getTimestamp(index, tsType.getPrecision()); + } else if (logicalType instanceof LocalZonedTimestampType) { + LocalZonedTimestampType ltzType = (LocalZonedTimestampType) logicalType; + return rowData.getTimestamp(index, ltzType.getPrecision()); + } else if (logicalType instanceof DecimalType) { + DecimalType decimalType = (DecimalType) logicalType; + return rowData.getDecimal(index, decimalType.getPrecision(), decimalType.getScale()); } else if (logicalType instanceof ArrayType) { return rowData.getArray(index); } else if (logicalType instanceof RowType) { @@ -422,8 +500,14 @@ private void writeValue(FieldVector vector, int index, Object value, LogicalType ((FixedSizeBinaryVector) vector).setSafe(index, (byte[]) value); } else if (logicalType instanceof DateType) { ((DateDayVector) vector).setSafe(index, (int) value); + } else if (logicalType instanceof TimeType) { + writeTime(vector, index, (int) value); } else if (logicalType instanceof TimestampType) { writeTimestamp(vector, index, (TimestampData) value, (TimestampType) logicalType); + } else if (logicalType instanceof LocalZonedTimestampType) { + writeLocalZonedTimestamp(vector, index, (TimestampData) value); + } else if (logicalType instanceof DecimalType) { + ((DecimalVector) vector).setSafe(index, ((DecimalData) value).toBigDecimal()); } else if (logicalType instanceof ArrayType) { writeArray(vector, index, (ArrayData) value, (ArrayType) logicalType); } else if (logicalType instanceof RowType) { @@ -468,6 +552,24 @@ private void setNull(FieldVector vector, int index) { ((TimeStampMicroVector) vector).setNull(index); } else if (vector instanceof TimeStampNanoVector) { ((TimeStampNanoVector) vector).setNull(index); + } else if (vector instanceof TimeStampSecTZVector) { + ((TimeStampSecTZVector) vector).setNull(index); + } else if (vector instanceof TimeStampMilliTZVector) { + ((TimeStampMilliTZVector) vector).setNull(index); + } else if (vector instanceof TimeStampMicroTZVector) { + ((TimeStampMicroTZVector) vector).setNull(index); + } else if (vector instanceof TimeStampNanoTZVector) { + ((TimeStampNanoTZVector) vector).setNull(index); + } else if (vector instanceof TimeSecVector) { + ((TimeSecVector) vector).setNull(index); + } else if (vector instanceof TimeMilliVector) { + ((TimeMilliVector) vector).setNull(index); + } else if (vector instanceof TimeMicroVector) { + ((TimeMicroVector) vector).setNull(index); + } else if (vector instanceof TimeNanoVector) { + ((TimeNanoVector) vector).setNull(index); + } else if (vector instanceof DecimalVector) { + ((DecimalVector) vector).setNull(index); } else if (vector instanceof FixedSizeListVector) { ((FixedSizeListVector) vector).setNull(index); } else if (vector instanceof ListVector) { @@ -500,6 +602,51 @@ private void writeTimestamp(FieldVector vector, int index, TimestampData tsData, } } + /** + * Write a TIME value given as milliseconds since midnight. + * + *

    The incoming value is always millisecond-resolution because that is Flink's internal + * representation, so the finer Arrow units are scaled up rather than truncated. + */ + private void writeTime(FieldVector vector, int index, int millisOfDay) { + if (vector instanceof TimeSecVector) { + ((TimeSecVector) vector).setSafe(index, millisOfDay / 1000); + } else if (vector instanceof TimeMilliVector) { + ((TimeMilliVector) vector).setSafe(index, millisOfDay); + } else if (vector instanceof TimeMicroVector) { + ((TimeMicroVector) vector).setSafe(index, millisOfDay * 1000L); + } else if (vector instanceof TimeNanoVector) { + ((TimeNanoVector) vector).setSafe(index, millisOfDay * 1_000_000L); + } else { + throw new LanceTypeConverter.UnsupportedTypeException( + "Unsupported time Vector type: " + vector.getClass().getSimpleName()); + } + } + + /** + * Write a TIMESTAMP_LTZ value into one of Arrow's zoned timestamp vectors. + * + *

    {@link TimestampData} already holds an epoch-based instant for this type, so the value is + * written as-is; applying a zone offset here would shift the instant. + */ + private void writeLocalZonedTimestamp(FieldVector vector, int index, TimestampData tsData) { + long millis = tsData.getMillisecond(); + int nanos = tsData.getNanoOfMillisecond(); + + if (vector instanceof TimeStampSecTZVector) { + ((TimeStampSecTZVector) vector).setSafe(index, millis / 1000); + } else if (vector instanceof TimeStampMilliTZVector) { + ((TimeStampMilliTZVector) vector).setSafe(index, millis); + } else if (vector instanceof TimeStampMicroTZVector) { + ((TimeStampMicroTZVector) vector).setSafe(index, millis * 1000L + nanos / 1000); + } else if (vector instanceof TimeStampNanoTZVector) { + ((TimeStampNanoTZVector) vector).setSafe(index, millis * 1_000_000L + nanos); + } else { + throw new LanceTypeConverter.UnsupportedTypeException( + "Unsupported zoned timestamp Vector type: " + vector.getClass().getSimpleName()); + } + } + /** * Write array value */ diff --git a/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterNewTypesTest.java b/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterNewTypesTest.java new file mode 100644 index 0000000..9ca5a27 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterNewTypesTest.java @@ -0,0 +1,234 @@ +/* + * 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.converter; + +import org.apache.flink.table.data.DecimalData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.TimestampData; +import org.apache.flink.table.types.logical.DecimalType; +import org.apache.flink.table.types.logical.LocalZonedTimestampType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.TimeType; +import org.apache.flink.table.types.logical.TimestampType; + +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.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.math.BigDecimal; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Covers DECIMAL, TIME and TIMESTAMP_LTZ, which had no Arrow mapping at all: {@code + * flinkTypeToArrowField} rejected them outright, so such a table could not be created even though + * {@code RowDataFieldAccessor} already encoded these types as primary keys. + * + *

    Assertions run over a full write-then-read cycle rather than over the schema mapping alone, + * because a type can map cleanly and still lose its value in the vector round-trip. + */ +class LanceTypeConverterNewTypesTest { + + private BufferAllocator allocator; + + @BeforeEach + void setUp() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + void tearDown() { + allocator.close(); + } + + /** Writes {@code rows} through the converter and reads them back out. */ + private List roundTrip(RowType rowType, List rows) { + RowDataConverter converter = new RowDataConverter(rowType); + try (VectorSchemaRoot root = converter.createVectorSchemaRoot(allocator)) { + converter.toVectorSchemaRoot(rows, root); + return converter.toRowDataList(root); + } + } + + private static RowType rowTypeOf(String name, LogicalType type) { + return new RowType(Collections.singletonList(new RowType.RowField(name, type))); + } + + // ------------------------------------------------------------------ + // schema mapping + // ------------------------------------------------------------------ + + @Test + @DisplayName("DECIMAL maps to Arrow Decimal preserving precision and scale") + void decimalMapsToArrowDecimal() { + Field field = + LanceTypeConverter.flinkTypeToArrowField("amount", new DecimalType(true, 10, 2)); + + assertThat(field.getType()).isInstanceOf(ArrowType.Decimal.class); + ArrowType.Decimal decimal = (ArrowType.Decimal) field.getType(); + assertThat(decimal.getPrecision()).isEqualTo(10); + assertThat(decimal.getScale()).isEqualTo(2); + } + + @Test + @DisplayName("TIME picks the Arrow bit width its unit requires") + void timeUsesUnitCompatibleBitWidth() { + // Arrow rejects Time32 paired with MICROSECOND and Time64 paired with MILLISECOND, so the + // width has to track the precision rather than being fixed. + ArrowType milli = LanceTypeConverter.flinkTypeToArrowField("a", new TimeType(true, 3)).getType(); + ArrowType micro = LanceTypeConverter.flinkTypeToArrowField("b", new TimeType(true, 6)).getType(); + + assertThat(((ArrowType.Time) milli).getBitWidth()).isEqualTo(32); + assertThat(((ArrowType.Time) micro).getBitWidth()).isEqualTo(64); + } + + @Test + @DisplayName("TIMESTAMP_LTZ is tagged with a timezone, plain TIMESTAMP is not") + void zonedAndUnzonedTimestampsStayDistinct() { + ArrowType zoned = + LanceTypeConverter.flinkTypeToArrowField("a", new LocalZonedTimestampType(true, 6)) + .getType(); + ArrowType unzoned = + LanceTypeConverter.flinkTypeToArrowField("b", new TimestampType(true, 6)).getType(); + + assertThat(((ArrowType.Timestamp) zoned).getTimezone()).isEqualTo("UTC"); + assertThat(((ArrowType.Timestamp) unzoned).getTimezone()) + .as("an unzoned TIMESTAMP must not acquire a zone") + .isNull(); + } + + @Test + @DisplayName("Arrow zoned timestamp reads back as TIMESTAMP_LTZ, not TIMESTAMP") + void zonedTimestampReadsBackAsLtz() { + // Without the timezone check on the reverse mapping both Arrow forms would collapse into + // TIMESTAMP and the round-trip would silently change the column's type. + Schema schema = + new Schema( + Arrays.asList( + LanceTypeConverter.flinkTypeToArrowField( + "zoned", new LocalZonedTimestampType(true, 6)), + LanceTypeConverter.flinkTypeToArrowField( + "plain", new TimestampType(true, 6)))); + + RowType recovered = LanceTypeConverter.toFlinkRowType(schema); + + assertThat(recovered.getTypeAt(0)).isInstanceOf(LocalZonedTimestampType.class); + assertThat(recovered.getTypeAt(1)).isInstanceOf(TimestampType.class); + } + + // ------------------------------------------------------------------ + // value round-trip + // ------------------------------------------------------------------ + + @Test + @DisplayName("DECIMAL values survive the vector round-trip") + void decimalValueRoundTrips() { + RowType rowType = rowTypeOf("amount", new DecimalType(true, 10, 2)); + DecimalData value = DecimalData.fromBigDecimal(new BigDecimal("123.45"), 10, 2); + + List out = roundTrip(rowType, Collections.singletonList(GenericRowData.of(value))); + + assertThat(out).hasSize(1); + assertThat(out.get(0).getDecimal(0, 10, 2).toBigDecimal()) + .isEqualByComparingTo(new BigDecimal("123.45")); + } + + @Test + @DisplayName("Negative DECIMAL values keep their sign") + void negativeDecimalRoundTrips() { + RowType rowType = rowTypeOf("amount", new DecimalType(true, 10, 2)); + DecimalData value = DecimalData.fromBigDecimal(new BigDecimal("-7.89"), 10, 2); + + List out = roundTrip(rowType, Collections.singletonList(GenericRowData.of(value))); + + assertThat(out.get(0).getDecimal(0, 10, 2).toBigDecimal()) + .isEqualByComparingTo(new BigDecimal("-7.89")); + } + + @Test + @DisplayName("TIME values survive the vector round-trip") + void timeValueRoundTrips() { + RowType rowType = rowTypeOf("at", new TimeType(true, 3)); + int millisOfDay = 45_296_123; // 12:34:56.123 + + List out = + roundTrip(rowType, Collections.singletonList(GenericRowData.of(millisOfDay))); + + assertThat(out.get(0).getInt(0)).isEqualTo(millisOfDay); + } + + @Test + @DisplayName("TIMESTAMP_LTZ values survive the vector round-trip") + void localZonedTimestampRoundTrips() { + RowType rowType = rowTypeOf("ts", new LocalZonedTimestampType(true, 6)); + TimestampData value = TimestampData.fromEpochMillis(1_700_000_000_123L); + + List out = roundTrip(rowType, Collections.singletonList(GenericRowData.of(value))); + + assertThat(out.get(0).getTimestamp(0, 6).getMillisecond()) + .isEqualTo(1_700_000_000_123L); + } + + @Test + @DisplayName("Pre-epoch TIMESTAMP_LTZ values are not corrupted by the micro split") + void preEpochLocalZonedTimestampRoundTrips() { + // Splitting epoch micros with % and / truncates towards zero, which yields a negative + // nanosecond remainder for instants before 1970; floorDiv/floorMod are used instead. + RowType rowType = rowTypeOf("ts", new LocalZonedTimestampType(true, 6)); + TimestampData value = TimestampData.fromEpochMillis(-1_000L); + + List out = roundTrip(rowType, Collections.singletonList(GenericRowData.of(value))); + + assertThat(out.get(0).getTimestamp(0, 6).getMillisecond()).isEqualTo(-1_000L); + } + + @Test + @DisplayName("NULLs in the new types are preserved rather than written as zero") + void nullsArePreserved() { + // setNull has no trailing else, so a vector type missing from its dispatch chain leaves the + // slot at its default value instead of marking it null. + RowType rowType = + new RowType( + Arrays.asList( + new RowType.RowField("amount", new DecimalType(true, 10, 2)), + new RowType.RowField("at", new TimeType(true, 3)), + new RowType.RowField("ts", new LocalZonedTimestampType(true, 6)))); + + List out = + roundTrip( + rowType, + Collections.singletonList(GenericRowData.of(null, null, null))); + + assertThat(out.get(0).isNullAt(0)).as("DECIMAL null").isTrue(); + assertThat(out.get(0).isNullAt(1)).as("TIME null").isTrue(); + assertThat(out.get(0).isNullAt(2)).as("TIMESTAMP_LTZ null").isTrue(); + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java b/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java index f18338f..6ac5bfc 100644 --- a/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java +++ b/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java @@ -166,13 +166,18 @@ void testUnresolvedDataTypeRejected() { @Test @DisplayName("An unsupported column type is rejected as a schema problem") void testUnsupportedTypeRejected() { + // MAP has no Arrow mapping yet. DECIMAL used to stand in here, so this case had to move + // when DECIMAL gained one; the assertion is about how an unmappable type is surfaced, not + // about MAP specifically. CatalogTable table = tableWith( - Schema.newBuilder().column("amount", DataTypes.DECIMAL(10, 2)).build()); + Schema.newBuilder() + .column("attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.INT())) + .build()); assertThatThrownBy(() -> LanceNamespaceCatalog.toArrowIpcSchema(table, allocator)) .isInstanceOf(org.apache.flink.table.catalog.exceptions.CatalogException.class) .hasMessageContaining("Cannot create a Lance table with this schema") - .hasMessageContaining("DecimalType"); + .hasMessageContaining("MapType"); } @Test From c94b6c6bd27758061674993f51533821e155e640 Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 12:54:00 +0800 Subject: [PATCH 05/11] feat: support UPDATE via SupportsRowLevelUpdate The sink now implements SupportsRowLevelUpdate, so UPDATE ... WHERE runs against a keyed Lance table. UPDATED_ROWS is requested because that mode delivers only the matched rows, each tagged UPDATE_AFTER, which is what the existing keyed upsert path already consumes; ALL_ROWS would stream back rows the statement never touched and rewrite them for no gain. requiredColumns() returns empty, meaning every column. The sink writes whole rows through mergeInsert(withMatchedUpdateAll), so narrowing the projection to the SET list plus the key would null out every unmentioned column. UPDATE on a table without a primary key is rejected during planning, since there is nothing for mergeInsert to match on. Two pre-existing defects blocked the end-to-end test and are fixed here. A keyed write lost everything it buffered when a job ended without a checkpoint. close() deliberately does not flush, because it also runs on cancel and failover, but nothing covered a graceful end of input -- so a batch job, which never checkpoints, completed reporting success while writing nothing. This is now handled in finish(), which Flink calls only on the normal completion path, preserving the checkpoint-as-persistence-boundary contract. The append path was unaffected because LanceSink.close() does flush, so the two sinks disagreed on whether a completed bounded job was durable. createDynamicTableSource/Sink passed the hadoop.* prefix list straight to validateExcept, which rejects an empty array. Any table declaring no hadoop option therefore failed validation outright, which is every table not backed by HDFS. The prefix list is now checked before choosing the overload. Both new test classes are named *Test rather than *ITCase: surefire's default includes do not match the ITCase suffix and no failsafe execution is configured, so the project's 86 existing ITCase cases never ran under mvn test. That gap and a long-standing concurrent-first-write failure it was hiding are recorded in .gh-comments/issue-update-followups.md; neither is in scope here. --- .gh-comments/issue-update-followups.md | 77 ++++++ lance-flink-docs-ability-map.md | 6 +- .../connector/lance/LanceUpsertSink.java | 16 +- .../lance/table/LanceDynamicTableFactory.java | 21 +- .../lance/table/LanceDynamicTableSink.java | 43 +++- .../lance/table/LanceRowLevelUpdateTest.java | 222 ++++++++++++++++++ .../table/LanceUpsertSinkFinishTest.java | 146 ++++++++++++ 7 files changed, 523 insertions(+), 8 deletions(-) create mode 100644 .gh-comments/issue-update-followups.md create mode 100644 src/test/java/org/apache/flink/connector/lance/table/LanceRowLevelUpdateTest.java create mode 100644 src/test/java/org/apache/flink/connector/lance/table/LanceUpsertSinkFinishTest.java diff --git a/.gh-comments/issue-update-followups.md b/.gh-comments/issue-update-followups.md new file mode 100644 index 0000000..7cdb5ae --- /dev/null +++ b/.gh-comments/issue-update-followups.md @@ -0,0 +1,77 @@ +# 本轮（UPDATE 实现）过程中暴露的既有缺陷 + +这两项都不是 `UPDATE` 引入的，也不在其修复范围内。它们是实现 `UPDATE` 时因首次有端到端 +SQL 测试覆盖而暴露出来的既有问题，单独记录以免随实现一起被淡化。 + +--- + +## C1：`*ITCase` 测试在 `mvn test` 中从未被执行 + +**现象** + +surefire 默认只匹配 `*Test` / `Test*` / `*Tests` / `*TestCase`，`*ITCase` 不在其中， +而项目未配置 failsafe 插件、也未自定义 ``。因此以下测试类从未在全量构建中运行： + +``` +LanceUpsertSinkITCase 9 个用例（含 1 个长期失败，见 C2） +LanceCatalogTableChangeITCase 8 +LanceNamespaceCatalogITCase ... +LanceCatalogTableITCase ... +LanceTimeTravelITCase 4 +LanceSinkConcurrencyITCase 1 +CompositePkDeleteITCase 1 +``` + +单独执行 `-Dtest='*ITCase'` 时共 86 个用例。 + +**影响** + +此前多轮报告的「全量 N 个测试通过」均未包含这 86 个用例，其中含一个真实失败。 +A6 的 `LanceCatalogTableChangeITCase` 等验收测试实际上从未在全量回归中生效。 + +**处置** + +本轮新增的两个测试类已命名为 `*Test` 以确保执行。既有 ITCase 未改名，因为重命名会 +影响 CI 中可能存在的分阶段执行约定（如故意把集成测试留给单独的 job）。 + +**待决策** + +需要确认 CI 是否另有 `-Dtest='*ITCase'` 阶段。若没有，应二者择一： +配置 failsafe 并绑定 `verify`，或把 ITCase 纳入 surefire 的 ``。 +在此之前，「全量通过」的说法应明确排除 ITCase。 + +--- + +## C2：两个 subtask 并发首写时主键元数据提交冲突 + +**现象** + +`LanceUpsertSinkITCase#twoSubtasksConcurrentFirstWrite` 长期失败： + +``` +RuntimeException: Incompatible transaction: This UpdateConfig transaction is +incompatible with concurrent transaction UpdateConfig at version 2. + at PrimaryKeyPersistence.persist(PrimaryKeyPersistence.java:73) + at LanceUpsertSink.open(LanceUpsertSink.java:211) +``` + +已确认为既有缺陷：在本轮改动前的提交上 stash 验证，同样失败。 + +**成因** + +`LanceUpsertSink.open()` 对首写执行 open-or-create，随后每个 subtask 都调用 +`PrimaryKeyPersistence.persist` 写入主键元数据。该写入是一个 `UpdateConfig` 事务， +多个 subtask 并发提交时 Lance 的冲突解析器直接拒绝，而非合并。 + +**影响** + +并行度 > 1 的表在首次写入时可能整体失败。dataset 已存在时不受影响， +因此在先建表再写入的流程中不易触发。 + +**候选方向** + +- 仅由 subtask 0 写元数据，其余 subtask 跳过 +- `persist` 捕获冲突异常后重读校验，若目标值已一致则视为成功 +- 把主键元数据的写入前移到 catalog 建表阶段，使 sink 的 open 不再写 config + +需要先确认 Lance 是否为 `UpdateConfig` 提供幂等或可重试的提交语义，再选方案。 diff --git a/lance-flink-docs-ability-map.md b/lance-flink-docs-ability-map.md index 703c7af..050341d 100644 --- a/lance-flink-docs-ability-map.md +++ b/lance-flink-docs-ability-map.md @@ -23,7 +23,7 @@ | `config.md` | 全部 ConfigOption（读/写/索引/向量/S3/hadoop.*） | — | | `operations/ddl/` | CREATE TABLE、CREATE CATALOG、database 操作 | ALTER TABLE、CREATE INDEX、分区 | | `operations/dql/` | SELECT（列裁剪/过滤/limit/聚合下推）、向量搜索、时间旅行 | 全文搜索、混合搜索 | -| `operations/dml/` | INSERT INTO（append）、INSERT OVERWRITE | UPDATE、DELETE、主键/upsert | +| `operations/dml/` | INSERT INTO（append）、INSERT OVERWRITE、UPDATE、DELETE、主键/upsert | — | | `performance.md` | 索引类型、向量调优参数 | benchmark 数据 | ## 3. 各文件详细内容清单 @@ -99,7 +99,7 @@ |---|---|---| | `INSERT INTO`（append） | ✅ | `write.mode=append`（默认） | | `INSERT OVERWRITE` | ✅ | `write.mode=overwrite`（首次写或 overwrite 模式） | -| `UPDATE` | ❌ | 未实现 | +| `UPDATE` | ✅ | `SupportsRowLevelUpdate`，`UPDATED_ROWS` 模式经主键 `mergeInsert` upsert；要求声明主键 | | `DELETE` | ✅ | 经主键 key-only `mergeInsert` + `withMatchedDelete`（PR #76 + follow-up A4） | | 主键 / upsert | ✅ | `PRIMARY KEY NOT ENFORCED` 时声明 `+I/+U/-D`，`keyBy` 保序后走 `mergeInsert` | @@ -117,7 +117,7 @@ | SQL 语法定位 | 交互式 DDL/DML | 声明式长驻拓扑（CREATE TABLE + INSERT INTO） | | vendors | databricks | 无（tbdsfs/hdfs 通过 `hadoop.*` 前缀支撑） | | 全文/混合搜索 | ✅ | ❌ | -| DELETE/UPDATE | ✅ | ❌（进行中） | +| DELETE/UPDATE | ✅ | ✅（需声明主键） | ## 5. 建议的文档落地顺序 diff --git a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java index fd1bcd2..ae7f29c 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java @@ -91,7 +91,8 @@ *

    Consistency model

    *

    This sink provides at-least-once semantics, not exactly-once: *

    * *

    Concurrency and first-write

    @@ -466,6 +469,15 @@ public void initializeState(FunctionInitializationContext context) { LOG.debug("Initialize state, isRestored: {}", context.isRestored()); } + @Override + public void finish() throws Exception { + // Called once the input is exhausted and only on the normal completion path, unlike + // close(), which also runs on cancel and failover. A batch job has no checkpoint, so + // without this the buffered rows of a keyed write would never reach the dataset. + LOG.info("Input finished, flushing remaining {} buffered key(s)", buffer.size()); + flush(); + } + @Override public void close() throws Exception { LOG.info("Closing Lance Upsert Sink"); 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 21cb849..90ec0e9 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 @@ -196,7 +196,7 @@ public Set> optionalOptions() { public DynamicTableSource createDynamicTableSource(Context context) { FactoryUtil.TableFactoryHelper helper = FactoryUtil.createTableFactoryHelper(this, context); Map tableOptions = context.getCatalogTable().getOptions(); - helper.validateExcept(extractHadoopOptionKeys(tableOptions)); + validateAllowingHadoopOptions(helper, tableOptions); ReadableConfig config = helper.getOptions(); LanceOptions options = buildLanceOptions(config, tableOptions); @@ -211,7 +211,7 @@ public DynamicTableSource createDynamicTableSource(Context context) { public DynamicTableSink createDynamicTableSink(Context context) { FactoryUtil.TableFactoryHelper helper = FactoryUtil.createTableFactoryHelper(this, context); Map tableOptions = context.getCatalogTable().getOptions(); - helper.validateExcept(extractHadoopOptionKeys(tableOptions)); + validateAllowingHadoopOptions(helper, tableOptions); ReadableConfig config = helper.getOptions(); LanceOptions options = buildLanceOptions(config, tableOptions); @@ -259,6 +259,23 @@ private int[] resolvePrimaryKeyIndices(ResolvedSchema schema, List prima /** * 提取以 {@code hadoop.} 为前缀的选项 key，供 {@code validateExcept} 跳过校验。 */ + /** + * Validate table options, tolerating the freeform {@code hadoop.*} passthrough keys. + * + *

    {@code validateExcept} rejects an empty prefix array outright, so a table that declares no + * hadoop option cannot go through that overload at all — which is every table that does not + * target HDFS. The prefix list has to be checked before choosing which overload to call. + */ + private void validateAllowingHadoopOptions( + FactoryUtil.TableFactoryHelper helper, Map tableOptions) { + String[] hadoopKeys = extractHadoopOptionKeys(tableOptions); + if (hadoopKeys.length == 0) { + helper.validate(); + } else { + helper.validateExcept(hadoopKeys); + } + } + private String[] extractHadoopOptionKeys(Map tableOptions) { if (tableOptions == null) { return new String[0]; 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 c012358..3972c4e 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 @@ -25,26 +25,36 @@ 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.catalog.Column; import org.apache.flink.table.connector.ChangelogMode; import org.apache.flink.table.connector.ProviderContext; +import org.apache.flink.table.connector.RowLevelModificationScanContext; 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.connector.sink.abilities.SupportsRowLevelUpdate; 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 javax.annotation.Nullable; + import java.util.Collections; import java.util.List; +import java.util.Optional; /** * Lance dynamic table sink. * *

    Implements DynamicTableSink interface, supports writing Flink data to Lance dataset. + * + *

    With a primary key declared, the sink also serves {@code UPDATE} through {@link + * SupportsRowLevelUpdate}. The statement is answered by the same keyed upsert path used for + * streaming writes, because Lance's {@code mergeInsert} already expresses match-by-key-then-update. */ -public class LanceDynamicTableSink implements DynamicTableSink { +public class LanceDynamicTableSink implements DynamicTableSink, SupportsRowLevelUpdate { private final LanceOptions options; private final DataType physicalDataType; @@ -125,6 +135,37 @@ public DynamicTableSink copy() { return new LanceDynamicTableSink(options, physicalDataType, primaryKeys, primaryKeyIndices); } + @Override + public RowLevelUpdateInfo applyRowLevelUpdate( + List updatedColumns, @Nullable RowLevelModificationScanContext context) { + if (primaryKeys.isEmpty()) { + // Without a key there is nothing for mergeInsert to match on, so an update could only + // be served by rewriting the table. Fail during planning rather than at runtime. + throw new UnsupportedOperationException( + "UPDATE requires a PRIMARY KEY NOT ENFORCED on the Lance table. " + + "The table is append-only without one; " + + "declare a primary key to enable row-level updates."); + } + + return new RowLevelUpdateInfo() { + @Override + public Optional> requiredColumns() { + // Empty means "every column, in table order". The sink writes the full row through + // mergeInsert(withMatchedUpdateAll), so a projection limited to the SET list plus + // the key would null out every column the statement did not mention. + return Optional.empty(); + } + + @Override + public RowLevelUpdateMode getRowLevelUpdateMode() { + // UPDATED_ROWS delivers only the matched rows, each tagged UPDATE_AFTER, which is + // what the keyed upsert path already consumes. ALL_ROWS would stream back rows the + // statement did not touch and rewrite them for no gain. + return RowLevelUpdateMode.UPDATED_ROWS; + } + }; + } + @Override public String asSummaryString() { return "Lance Table Sink"; diff --git a/src/test/java/org/apache/flink/connector/lance/table/LanceRowLevelUpdateTest.java b/src/test/java/org/apache/flink/connector/lance/table/LanceRowLevelUpdateTest.java new file mode 100644 index 0000000..ec1eac9 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/table/LanceRowLevelUpdateTest.java @@ -0,0 +1,222 @@ +/* + * 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.EnvironmentSettings; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; +import org.apache.flink.types.Row; + +import org.apache.arrow.memory.RootAllocator; +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.ArrayList; +import java.util.Iterator; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * End-to-end coverage for {@code UPDATE}, served through {@link + * org.apache.flink.table.connector.sink.abilities.SupportsRowLevelUpdate}. + * + *

    Assertions read the data back rather than inspecting the returned {@code RowLevelUpdateInfo}, + * because the risk in this path is not whether the interface is wired up but whether the rewritten + * statement preserves columns the user never mentioned. + */ +class LanceRowLevelUpdateTest { + + @TempDir Path tempDir; + + private TableEnvironment tableEnv; + + @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() { + EnvironmentSettings settings = EnvironmentSettings.newInstance().inBatchMode().build(); + tableEnv = TableEnvironment.create(settings); + } + + private String datasetPath(String name) { + return tempDir.resolve(name + ".lance").toString(); + } + + /** Creates a keyed table and seeds it with three rows. */ + private void createSeededTable(String table) throws Exception { + tableEnv.executeSql( + "CREATE TABLE " + + table + + " (" + + " id BIGINT NOT NULL," + + " name STRING," + + " score DOUBLE," + + " PRIMARY KEY (id) NOT ENFORCED" + + ") WITH (" + + " 'connector' = 'lance'," + + " 'path' = '" + + datasetPath(table) + + "'" + + ")"); + + tableEnv.executeSql( + "INSERT INTO " + + table + + " VALUES " + + "(1, 'alice', 10.0), (2, 'bob', 20.0), (3, 'carol', 30.0)") + .await(); + } + + private List query(String sql) throws Exception { + List rows = new ArrayList<>(); + TableResult result = tableEnv.executeSql(sql); + try (org.apache.flink.util.CloseableIterator it = result.collect()) { + while (it.hasNext()) { + rows.add(it.next()); + } + } + return rows; + } + + private Row rowWithId(List rows, long id) { + return rows.stream() + .filter(r -> id == (Long) r.getField(0)) + .findFirst() + .orElseThrow(() -> new AssertionError("no row with id=" + id)); + } + + @Test + @DisplayName("UPDATE changes only the rows matching WHERE") + void updateAffectsOnlyMatchingRows() throws Exception { + createSeededTable("t_update_basic"); + + tableEnv.executeSql("UPDATE t_update_basic SET score = 99.0 WHERE id = 2").await(); + + List rows = query("SELECT id, name, score FROM t_update_basic"); + assertThat(rows).as("UPDATE must not change the row count").hasSize(3); + assertThat(rowWithId(rows, 2).getField(2)).isEqualTo(99.0); + assertThat(rowWithId(rows, 1).getField(2)).isEqualTo(10.0); + assertThat(rowWithId(rows, 3).getField(2)).isEqualTo(30.0); + } + + @Test + @DisplayName("Columns absent from SET keep their value instead of being nulled") + void untouchedColumnsSurvive() throws Exception { + // mergeInsert runs withMatchedUpdateAll, so it overwrites every column of a matched row. + // If requiredColumns() narrowed the projection to the SET list plus the key, every other + // column would arrive empty and be written as NULL. + createSeededTable("t_update_projection"); + + tableEnv.executeSql("UPDATE t_update_projection SET score = 55.0 WHERE id = 1").await(); + + Row updated = rowWithId(query("SELECT id, name, score FROM t_update_projection"), 1); + assertThat(updated.getField(1)).as("name was not in the SET list").isEqualTo("alice"); + assertThat(updated.getField(2)).isEqualTo(55.0); + } + + @Test + @DisplayName("UPDATE can set several columns at once") + void multiColumnUpdate() throws Exception { + createSeededTable("t_update_multi"); + + tableEnv.executeSql( + "UPDATE t_update_multi SET name = 'robert', score = 21.5 WHERE id = 2") + .await(); + + Row updated = rowWithId(query("SELECT id, name, score FROM t_update_multi"), 2); + assertThat(updated.getField(1)).isEqualTo("robert"); + assertThat(updated.getField(2)).isEqualTo(21.5); + } + + @Test + @DisplayName("An UPDATE matching several rows applies to all of them") + void updateMatchingManyRows() throws Exception { + createSeededTable("t_update_many"); + + tableEnv.executeSql("UPDATE t_update_many SET score = 0.0 WHERE score >= 20.0").await(); + + List rows = query("SELECT id, name, score FROM t_update_many"); + assertThat(rows).hasSize(3); + assertThat(rowWithId(rows, 1).getField(2)).as("below the filter").isEqualTo(10.0); + assertThat(rowWithId(rows, 2).getField(2)).isEqualTo(0.0); + assertThat(rowWithId(rows, 3).getField(2)).isEqualTo(0.0); + } + + @Test + @DisplayName("An UPDATE matching nothing leaves the table unchanged") + void updateMatchingNothingIsNoOp() throws Exception { + createSeededTable("t_update_nomatch"); + + tableEnv.executeSql("UPDATE t_update_nomatch SET score = 1.0 WHERE id = 999").await(); + + List rows = query("SELECT id, name, score FROM t_update_nomatch"); + assertThat(rows).hasSize(3); + assertThat(rowWithId(rows, 1).getField(2)).isEqualTo(10.0); + assertThat(rowWithId(rows, 2).getField(2)).isEqualTo(20.0); + assertThat(rowWithId(rows, 3).getField(2)).isEqualTo(30.0); + } + + @Test + @DisplayName("UPDATE can write NULL into a nullable column") + void updateToNull() throws Exception { + createSeededTable("t_update_null"); + + tableEnv.executeSql("UPDATE t_update_null SET name = CAST(NULL AS STRING) WHERE id = 3") + .await(); + + Row updated = rowWithId(query("SELECT id, name, score FROM t_update_null"), 3); + assertThat(updated.getField(1)).isNull(); + assertThat(updated.getField(2)).as("score must be untouched").isEqualTo(30.0); + } + + @Test + @DisplayName("UPDATE on a table without a primary key is rejected with a usable message") + void updateWithoutPrimaryKeyIsRejected() { + // There is nothing for mergeInsert to match on, so this has to fail during planning rather + // than silently append or fail deep inside the sink at runtime. + tableEnv.executeSql( + "CREATE TABLE t_update_nokey (" + + " id BIGINT NOT NULL," + + " name STRING" + + ") WITH (" + + " 'connector' = 'lance'," + + " 'path' = '" + + datasetPath("t_update_nokey") + + "'" + + ")"); + + assertThatThrownBy( + () -> + tableEnv.executeSql( + "UPDATE t_update_nokey SET name = 'x' WHERE id = 1")) + .hasStackTraceContaining("PRIMARY KEY"); + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/table/LanceUpsertSinkFinishTest.java b/src/test/java/org/apache/flink/connector/lance/table/LanceUpsertSinkFinishTest.java new file mode 100644 index 0000000..a5513f1 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/table/LanceUpsertSinkFinishTest.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.api.EnvironmentSettings; +import org.apache.flink.table.api.TableEnvironment; +import org.apache.flink.table.api.TableResult; +import org.apache.flink.types.Row; + +import org.apache.arrow.memory.RootAllocator; +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.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Guards durability at end-of-input for the keyed write path. + * + *

    {@code LanceUpsertSink} treats the checkpoint as its persistence boundary and deliberately + * does not flush from {@code close()}, since {@code close()} also runs on cancel and failover. A + * batch job never checkpoints, so before {@code finish()} was implemented every buffered row of a + * keyed write was discarded and the write completed reporting success while writing nothing. + * + *

    The append path was unaffected because {@code LanceSink.close()} does flush, so the two sinks + * disagreed on whether a completed bounded job was durable. These cases pin the keyed path. + */ +class LanceUpsertSinkFinishTest { + + @TempDir Path tempDir; + + private TableEnvironment tableEnv; + + @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() { + EnvironmentSettings settings = EnvironmentSettings.newInstance().inBatchMode().build(); + tableEnv = TableEnvironment.create(settings); + } + + private List query(String sql) throws Exception { + List rows = new ArrayList<>(); + TableResult result = tableEnv.executeSql(sql); + try (org.apache.flink.util.CloseableIterator it = result.collect()) { + while (it.hasNext()) { + rows.add(it.next()); + } + } + return rows; + } + + private void createTable(String table, boolean withPrimaryKey) { + tableEnv.executeSql( + "CREATE TABLE " + + table + + " (" + + " id BIGINT NOT NULL," + + " name STRING" + + (withPrimaryKey ? ", PRIMARY KEY (id) NOT ENFORCED" : "") + + ") WITH (" + + " 'connector' = 'lance'," + + " 'path' = '" + + tempDir.resolve(table + ".lance") + + "'" + + ")"); + } + + @Test + @DisplayName("A keyed batch INSERT is durable once the job completes") + void keyedBatchInsertIsDurable() throws Exception { + // The row count here stays below write.batch-size, so nothing triggers a size-based flush + // and durability rests entirely on end-of-input handling. + createTable("t_finish_keyed", true); + + tableEnv.executeSql("INSERT INTO t_finish_keyed VALUES (1, 'a'), (2, 'b')").await(); + + assertThat(query("SELECT id, name FROM t_finish_keyed")) + .as("a completed batch job must not report success while writing nothing") + .hasSize(2); + } + + @Test + @DisplayName("The append path is durable too, so both sinks agree") + void appendBatchInsertIsDurable() throws Exception { + createTable("t_finish_append", false); + + tableEnv.executeSql("INSERT INTO t_finish_append VALUES (1, 'a'), (2, 'b')").await(); + + assertThat(query("SELECT id, name FROM t_finish_append")).hasSize(2); + } + + @Test + @DisplayName("Consecutive batch writes accumulate instead of overwriting") + void consecutiveKeyedWritesAccumulate() throws Exception { + createTable("t_finish_repeat", true); + + tableEnv.executeSql("INSERT INTO t_finish_repeat VALUES (1, 'a')").await(); + tableEnv.executeSql("INSERT INTO t_finish_repeat VALUES (2, 'b')").await(); + + assertThat(query("SELECT id, name FROM t_finish_repeat")) + .as("the second job must not lose the first job's row") + .hasSize(2); + } + + @Test + @DisplayName("Re-inserting a key updates it rather than duplicating it") + void reinsertingKeyUpserts() throws Exception { + createTable("t_finish_upsert", true); + + tableEnv.executeSql("INSERT INTO t_finish_upsert VALUES (1, 'first')").await(); + tableEnv.executeSql("INSERT INTO t_finish_upsert VALUES (1, 'second')").await(); + + List rows = query("SELECT id, name FROM t_finish_upsert"); + assertThat(rows).hasSize(1); + assertThat(rows.get(0).getField(1)).isEqualTo("second"); + } +} From 0ddc62f9ca4d56f153f064e171573af94142ca1a Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 13:22:27 +0800 Subject: [PATCH 06/11] build: run ITCase via failsafe, and fix the concurrent first-write it exposed The project's *ITCase integration tests never ran in a build. surefire's default includes do not match the ITCase suffix, and although CI runs mvn verify, no failsafe plugin was configured, so the verify phase bound only surefire. Seventy-five integration tests -- including one long-standing failure -- were silently skipped on every "green" build. Configure maven-failsafe-plugin and bind its integration-test and verify goals. The verify goal is the operative one: without it the reports are written but a failure does not fail the build. failsafe's default includes already cover **/*ITCase.java, so the existing classes keep their standard Flink-ecosystem names. The version is 3.3.0 rather than surefire's 3.1.2 because 3.1.2 of failsafe is not available from the internal mirror; the two plugins version independently. Turning the tests on surfaced a real defect. PrimaryKeyPersistence.persist wrote the primary-key metadata with an unconditional updateConfig, which is a versioned Lance transaction rather than an idempotent put. Two subtasks first-writing the same dataset committed it at once and the conflict resolver rejected one outright. persist is now genuinely idempotent: it compares the stored value first so the steady state opens no transaction at all, and on a conflict it advances the handle with checkoutLatest before re-reading -- the handle is pinned to the version seen at open() and would otherwise never observe the peer's commit -- accepting the result when the winner stored the same value and rethrowing only on a genuine disagreement. Removing the checkoutLatest call reproduces the original conflict, which confirms the stale handle was the root cause. The conflict arrives as a bare RuntimeException from the Rust resolver with no dedicated type, so the outcome is verified by re-reading rather than by matching the message. Both C1 and C2 in .gh-comments/issue-update-followups.md are now resolved. mvn verify runs surefire 276 + failsafe 75 per module, green across 1.18 / 1.19 / 1.20. --- .gh-comments/issue-update-followups.md | 86 +++++------ pom.xml | 27 ++++ .../connector/lance/LanceUpsertSink.java | 3 +- .../lance/PrimaryKeyPersistence.java | 42 +++++- .../PrimaryKeyPersistenceConcurrencyTest.java | 140 ++++++++++++++++++ 5 files changed, 248 insertions(+), 50 deletions(-) create mode 100644 src/test/java/org/apache/flink/connector/lance/PrimaryKeyPersistenceConcurrencyTest.java diff --git a/.gh-comments/issue-update-followups.md b/.gh-comments/issue-update-followups.md index 7cdb5ae..1c0ceea 100644 --- a/.gh-comments/issue-update-followups.md +++ b/.gh-comments/issue-update-followups.md @@ -1,77 +1,67 @@ -# 本轮（UPDATE 实现）过程中暴露的既有缺陷 +# UPDATE 实现过程中暴露的既有缺陷（已解决） -这两项都不是 `UPDATE` 引入的，也不在其修复范围内。它们是实现 `UPDATE` 时因首次有端到端 -SQL 测试覆盖而暴露出来的既有问题，单独记录以免随实现一起被淡化。 +这两项都不是 `UPDATE` 引入的。它们是实现 `UPDATE` 时因首次有端到端 SQL 测试覆盖而暴露 +出来的既有问题，均已在本轮修复。保留本文件作为背景记录。 --- -## C1：`*ITCase` 测试在 `mvn test` 中从未被执行 +## C1：`*ITCase` 测试在构建中从未被执行 —— 已修复 -**现象** +**曾经的现象** -surefire 默认只匹配 `*Test` / `Test*` / `*Tests` / `*TestCase`，`*ITCase` 不在其中， -而项目未配置 failsafe 插件、也未自定义 ``。因此以下测试类从未在全量构建中运行： +surefire 的默认 includes 不匹配 `*ITCase` 后缀，而项目未配置 failsafe 插件。CI 跑的是 +`mvn verify`，但 `verify` 阶段的插件链里只有 surefire，没有任何 failsafe，因此 75 个 +集成测试从未在构建中运行——CI 每次"通过"都不包含它们，其中还藏着一个真实失败（见 C2）。 -``` -LanceUpsertSinkITCase 9 个用例（含 1 个长期失败，见 C2） -LanceCatalogTableChangeITCase 8 -LanceNamespaceCatalogITCase ... -LanceCatalogTableITCase ... -LanceTimeTravelITCase 4 -LanceSinkConcurrencyITCase 1 -CompositePkDeleteITCase 1 -``` - -单独执行 `-Dtest='*ITCase'` 时共 86 个用例。 - -**影响** +**修复** -此前多轮报告的「全量 N 个测试通过」均未包含这 86 个用例，其中含一个真实失败。 -A6 的 `LanceCatalogTableChangeITCase` 等验收测试实际上从未在全量回归中生效。 +在 `pom.xml` 的 pluginManagement 声明 `maven-failsafe-plugin`（3.3.0，因内部镜像无 3.1.2）， +并在 build/plugins 绑定 `integration-test` + `verify` 两个 goal。verify goal 是关键： +没有它，报告会生成但构建仍然通过。 -**处置** +failsafe 默认 includes 含 `**/*ITCase.java`，正好匹配项目约定，无需自定义。既有 ITCase +未改名，保留 Flink 生态的标准命名。 -本轮新增的两个测试类已命名为 `*Test` 以确保执行。既有 ITCase 未改名，因为重命名会 -影响 CI 中可能存在的分阶段执行约定（如故意把集成测试留给单独的 job）。 +**验证** -**待决策** - -需要确认 CI 是否另有 `-Dtest='*ITCase'` 阶段。若没有，应二者择一： -配置 failsafe 并绑定 `verify`，或把 ITCase 纳入 surefire 的 ``。 -在此之前，「全量通过」的说法应明确排除 ITCase。 +`mvn verify` 现在每模块执行 surefire 276 + failsafe 75，三个 Flink 版本模块全绿。 +移除 verify goal 曾确认失败的集成测试不会中断构建，加回后恢复拦截。 --- -## C2：两个 subtask 并发首写时主键元数据提交冲突 +## C2：两个 subtask 并发首写时主键元数据提交冲突 —— 已修复 -**现象** +**曾经的现象** -`LanceUpsertSinkITCase#twoSubtasksConcurrentFirstWrite` 长期失败： +`LanceUpsertSinkITCase#twoSubtasksConcurrentFirstWrite` 失败： ``` RuntimeException: Incompatible transaction: This UpdateConfig transaction is incompatible with concurrent transaction UpdateConfig at version 2. - at PrimaryKeyPersistence.persist(PrimaryKeyPersistence.java:73) - at LanceUpsertSink.open(LanceUpsertSink.java:211) + at PrimaryKeyPersistence.persist(...) + at LanceUpsertSink.open(...) ``` -已确认为既有缺陷：在本轮改动前的提交上 stash 验证，同样失败。 - **成因** -`LanceUpsertSink.open()` 对首写执行 open-or-create，随后每个 subtask 都调用 -`PrimaryKeyPersistence.persist` 写入主键元数据。该写入是一个 `UpdateConfig` 事务， -多个 subtask 并发提交时 Lance 的冲突解析器直接拒绝，而非合并。 +`PrimaryKeyPersistence.persist` 无条件调用 `dataset.updateConfig` 写主键元数据。该调用 +是一个版本化的 Lance 事务，不是幂等 put；多个 subtask 并发首写时同时提交，Lance 的冲突 +解析器直接拒绝而非合并。原注释"Idempotent; safe to call on every open"是错误假设。 + +**修复** -**影响** +让 `persist` 真正幂等，两道防线： -并行度 > 1 的表在首次写入时可能整体失败。dataset 已存在时不受影响， -因此在先建表再写入的流程中不易触发。 +1. 写入前先比较——值已一致则完全不发起事务，稳态零写入； +2. 冲突时先 `dataset.checkoutLatest()` 把句柄推进到最新版本（`getConfig` 读的是 open 时 + 的旧快照，不推进就永远看不到对端提交），再重读校验；若目标值已达成则视为成功， + 否则照常抛出。 -**候选方向** +第 2 步的句柄推进是关键：移除它测试立即复现原冲突，证明句柄陈旧才是根因。冲突异常是 +Rust 层的裸 `RuntimeException` 无专用类型，因此用重读校验结果而非匹配消息来判断。 -- 仅由 subtask 0 写元数据，其余 subtask 跳过 -- `persist` 捕获冲突异常后重读校验，若目标值已一致则视为成功 -- 把主键元数据的写入前移到 catalog 建表阶段，使 sink 的 open 不再写 config +**验证** -需要先确认 Lance 是否为 `UpdateConfig` 提供幂等或可重试的提交语义，再选方案。 +`twoSubtasksConcurrentFirstWrite` 连续 3 次稳定通过；移除 `checkoutLatest` 立即复现原 +冲突（证明测试与诊断有效）。另加 `PrimaryKeyPersistenceConcurrencyTest`（5 个单元用例）， +其中一条专门钉住"容错不得吞掉真实分歧"——不同的 key list 仍必须写入。 diff --git a/pom.xml b/pom.xml index f3397ac..24ba54f 100644 --- a/pom.xml +++ b/pom.xml @@ -288,6 +288,18 @@ 3.1.2 + + + org.apache.maven.plugins + maven-failsafe-plugin + 3.3.0 + + org.codehaus.mojo @@ -481,6 +493,21 @@ org.apache.maven.plugins maven-surefire-plugin + + org.apache.maven.plugins + maven-failsafe-plugin + + + + + integration-test + verify + + + + org.apache.maven.plugins maven-enforcer-plugin diff --git a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java index ae7f29c..dfc945b 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java @@ -207,7 +207,8 @@ public void open(Configuration parameters) throws Exception { // Open-or-create: atomic in the Lance native layer, safe under concurrent first-writes. this.dataset = openOrCreate(datasetPath); - // Persist the primary keys into the dataset config. Idempotent; safe to call on every open. + // Persist the primary keys into the dataset config. Tolerates a peer subtask committing + // the same value concurrently; see PrimaryKeyPersistence#persist. PrimaryKeyPersistence.persist(dataset, primaryKeys); LOG.info("Lance Upsert Sink opened, primary keys: {}", primaryKeys); diff --git a/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java b/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java index c862e16..e394819 100644 --- a/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java +++ b/src/main/java/org/apache/flink/connector/lance/PrimaryKeyPersistence.java @@ -50,6 +50,15 @@ private PrimaryKeyPersistence() { /** * Persist the primary-key column names into the dataset config. * + *

    Safe to call concurrently from several subtasks opening the same dataset. {@code + * updateConfig} is a versioned Lance transaction rather than an idempotent put, so two + * subtasks committing it at the same time are rejected outright by the conflict resolver + * ("Incompatible transaction: This UpdateConfig transaction is incompatible with concurrent + * transaction UpdateConfig"). Two things keep that from surfacing: the value is compared + * first so that the steady state issues no transaction at all, and a losing commit advances + * its handle to the latest version and accepts the outcome when the winner already stored the + * same value. Only a genuine disagreement is propagated. + * * @param dataset the Lance dataset (must already be materialized) * @param primaryKeys ordered primary-key column names; empty/null writes nothing * @throws IllegalArgumentException if any column name contains a comma (which would make @@ -70,7 +79,38 @@ public static void persist(Dataset dataset, List primaryKeys) { + "comma-delimited encoding): '" + column + "'"); } } - dataset.updateConfig(Collections.singletonMap(PK_CONFIG_KEY, String.join(",", primaryKeys))); + + String encoded = String.join(",", primaryKeys); + if (encoded.equals(readRaw(dataset))) { + // Already stored, by an earlier open or by a peer subtask that got here first. + // Skipping keeps the steady state free of write transactions entirely. + return; + } + + try { + dataset.updateConfig(Collections.singletonMap(PK_CONFIG_KEY, encoded)); + } catch (RuntimeException e) { + // A peer may have committed the identical value between the check above and this + // commit. Lance surfaces that as a plain RuntimeException from the Rust conflict + // resolver with no dedicated type, so the outcome is verified by re-reading rather + // than by matching on the message. + // + // The handle has to be advanced first: it is pinned to the version observed at open() + // and getConfig() would keep returning the pre-conflict snapshot, making every peer + // commit look like a genuine disagreement. Advancing is harmless for the caller -- + // the sink wants the newest version anyway. + dataset.checkoutLatest(); + if (encoded.equals(readRaw(dataset))) { + return; + } + throw e; + } + } + + /** Returns the raw encoded primary-key config value, or {@code null} when absent. */ + private static String readRaw(Dataset dataset) { + Map config = dataset.getConfig(); + return config == null ? null : config.get(PK_CONFIG_KEY); } /** diff --git a/src/test/java/org/apache/flink/connector/lance/PrimaryKeyPersistenceConcurrencyTest.java b/src/test/java/org/apache/flink/connector/lance/PrimaryKeyPersistenceConcurrencyTest.java new file mode 100644 index 0000000..8e02379 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/PrimaryKeyPersistenceConcurrencyTest.java @@ -0,0 +1,140 @@ +/* + * 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.RootAllocator; +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.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.lance.Dataset; +import org.lance.WriteParams; + +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.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Covers the idempotence and conflict handling of {@link PrimaryKeyPersistence#persist}. + * + *

    {@code updateConfig} is a versioned Lance transaction, not an idempotent put, so several + * subtasks opening the same dataset used to be rejected by the conflict resolver. The tolerance + * added for that must not swallow a real disagreement, which is what the last case pins down. + */ +class PrimaryKeyPersistenceConcurrencyTest { + + @TempDir Path tempDir; + + private static final List KEYS = Arrays.asList("id", "region"); + + private String newDataset(String name) { + String path = tempDir.resolve(name + ".lance").toString(); + Schema schema = + new Schema( + Collections.singletonList( + new Field( + "id", + FieldType.nullable(new ArrowType.Int(64, true)), + null))); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Dataset.create(allocator, path, schema, new WriteParams.Builder().build()).close(); + } + return path; + } + + @Test + @DisplayName("persist writes the encoded key list on first call") + void persistWritesOnFirstCall() { + String path = newDataset("first"); + try (Dataset dataset = Dataset.open(path)) { + PrimaryKeyPersistence.persist(dataset, KEYS); + assertThat(PrimaryKeyPersistence.load(dataset)).isEqualTo(KEYS); + } + } + + @Test + @DisplayName("Repeating persist with the same value commits no further version") + void repeatedPersistIsAVersionNoOp() { + // The comparison before the write is what keeps the steady state free of transactions. + // Without it every subtask open would append a version to the dataset history. + String path = newDataset("repeat"); + try (Dataset dataset = Dataset.open(path)) { + PrimaryKeyPersistence.persist(dataset, KEYS); + long afterFirst = dataset.version(); + + PrimaryKeyPersistence.persist(dataset, KEYS); + + assertThat(dataset.version()) + .as("an unchanged value must not open a new transaction") + .isEqualTo(afterFirst); + } + } + + @Test + @DisplayName("A second handle observing the same value also stays quiet") + void peerHandleWithSameValueIsANoOp() { + // Models the common case: one subtask already stored the keys, another opens the dataset + // afterwards and finds them present. + String path = newDataset("peer"); + try (Dataset first = Dataset.open(path)) { + PrimaryKeyPersistence.persist(first, KEYS); + } + try (Dataset second = Dataset.open(path)) { + long before = second.version(); + PrimaryKeyPersistence.persist(second, KEYS); + assertThat(second.version()).isEqualTo(before); + assertThat(PrimaryKeyPersistence.load(second)).isEqualTo(KEYS); + } + } + + @Test + @DisplayName("A genuinely different key list is still written") + void differingValueIsWritten() { + String path = newDataset("differ"); + try (Dataset dataset = Dataset.open(path)) { + PrimaryKeyPersistence.persist(dataset, Collections.singletonList("id")); + PrimaryKeyPersistence.persist(dataset, KEYS); + assertThat(PrimaryKeyPersistence.load(dataset)) + .as("the conflict tolerance must not turn persist into a no-op") + .isEqualTo(KEYS); + } + } + + @Test + @DisplayName("A comma in a column name is rejected before any write") + void commaInColumnNameIsRejected() { + String path = newDataset("comma"); + try (Dataset dataset = Dataset.open(path)) { + assertThatThrownBy( + () -> + PrimaryKeyPersistence.persist( + dataset, Collections.singletonList("a,b"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("comma"); + assertThat(PrimaryKeyPersistence.load(dataset)).isEmpty(); + } + } +} From 7238592422a36e604b900c005fba9bdc6324c845 Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 14:15:55 +0800 Subject: [PATCH 07/11] fix: reject unsupported vectors in setNull instead of failing silently RowDataConverter#setNull ended its if-else chain without an else, so an Arrow vector type it did not enumerate was skipped without a word. That is data corruption rather than a harmless no-op: the validity bit of a slot that already holds a value stays set, so the previous row's value is serialised as this row's value and nothing is reported. A probe on IntervalDayVector confirmed it -- after setNull the slot still read back as non-null with the stale value intact. The bug stayed invisible because a freshly allocated vector masks it. The zeroed validity buffer reads back as null while getNullCount() still returns 0, so the two disagree and a first-write smoke test looks correct. Only a reused slot exposes the stale value. readValue, getFieldValue and writeValue all reject unknown types, so the silent branch was also an inconsistency inside a single class. setNull now throws UnsupportedTypeException naming the vector class and the field. MapVector is matched ahead of ListVector because it is a subclass; with the list branch first every map would be nulled through the list path. No Flink type maps to MapVector yet, but the ordering has to be right before MAP support lands, and it is cheaper to fix now than to debug later. The change is a net-new guard: the full suite passes unchanged, which confirms every type on the current production path is already enumerated. --- .../lance/converter/RowDataConverter.java | 16 +++ .../RowDataConverterSetNullFallbackTest.java | 132 ++++++++++++++++++ 2 files changed, 148 insertions(+) create mode 100644 src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterSetNullFallbackTest.java diff --git a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java index 992fcfd..a102bd4 100644 --- a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java +++ b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java @@ -73,6 +73,7 @@ import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.types.pojo.Schema; import org.slf4j.Logger; @@ -572,10 +573,25 @@ private void setNull(FieldVector vector, int index) { ((DecimalVector) vector).setNull(index); } else if (vector instanceof FixedSizeListVector) { ((FixedSizeListVector) vector).setNull(index); + } else if (vector instanceof MapVector) { + // Must precede ListVector: MapVector extends ListVector, so the ListVector branch + // would otherwise swallow it and null the map as if it were a plain list. + ((MapVector) vector).setNull(index); } else if (vector instanceof ListVector) { ((ListVector) vector).setNull(index); } else if (vector instanceof StructVector) { ((StructVector) vector).setNull(index); + } else { + // Falling through silently is data corruption, not a harmless no-op. The validity + // bit of a slot that already holds a value stays set, so the previous row's value is + // emitted as this row's value with no error anywhere. A freshly allocated vector + // hides it -- the zeroed validity buffer reads back as null while getNullCount() + // still reports 0 -- which is why this went unnoticed. readValue, getFieldValue and + // writeValue all reject unknown types; this branch makes setNull consistent. + throw new LanceTypeConverter.UnsupportedTypeException( + "Cannot write NULL: unsupported Arrow vector " + + vector.getClass().getSimpleName() + + " for field '" + vector.getField().getName() + "'"); } } diff --git a/src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterSetNullFallbackTest.java b/src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterSetNullFallbackTest.java new file mode 100644 index 0000000..83e9753 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterSetNullFallbackTest.java @@ -0,0 +1,132 @@ +/* + * 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.converter; + +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.IntervalDayVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.RowType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Constructor; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Pins the fallback branch of {@link RowDataConverter}'s {@code setNull}. + * + *

    The if-else chain used to end without an {@code else}, so an Arrow vector type it did not + * enumerate was skipped in silence. That is not a harmless no-op: the validity bit of a slot that + * already holds a value stays set, and the previous row's value is emitted as this row's value + * with nothing reported anywhere. A freshly allocated vector masks the bug, because the zeroed + * validity buffer reads back as null while {@code getNullCount()} still returns 0. + * + *

    {@code readValue}, {@code getFieldValue} and {@code writeValue} all reject unknown types, so + * the silent branch was also an inconsistency inside one class. + */ +class RowDataConverterSetNullFallbackTest { + + private static Method setNullMethod() throws Exception { + Method m = RowDataConverter.class.getDeclaredMethod("setNull", FieldVector.class, int.class); + m.setAccessible(true); + return m; + } + + private static RowDataConverter newConverter() throws Exception { + Constructor c = RowDataConverter.class.getDeclaredConstructor(RowType.class); + c.setAccessible(true); + return (RowDataConverter) c.newInstance(RowType.of(new IntType())); + } + + /** Unwraps the reflective wrapper so the assertion sees the real exception. */ + private static void invokeSetNull(FieldVector vector, int index) throws Exception { + try { + setNullMethod().invoke(newConverter(), vector, index); + } catch (InvocationTargetException e) { + if (e.getCause() instanceof Exception) { + throw (Exception) e.getCause(); + } + throw e; + } + } + + @Test + @DisplayName("An unsupported vector type is rejected instead of silently keeping a stale value") + void unsupportedVectorIsRejected() throws Exception { + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + // IntervalDayVector is deliberately outside the enumerated chain: no Flink + // logical type maps to it, so it stands in for any future gap. + IntervalDayVector vector = new IntervalDayVector("iv", allocator)) { + vector.allocateNew(2); + vector.setSafe(0, 1, 5_000); + vector.setValueCount(1); + + // Before the fix this call returned quietly and left slot 0 valid, so the row + // serialised with the stale value rather than NULL. + assertThatThrownBy(() -> invokeSetNull(vector, 0)) + .isInstanceOf(LanceTypeConverter.UnsupportedTypeException.class) + .hasMessageContaining("IntervalDayVector") + .hasMessageContaining("iv"); + } + } + + @Test + @DisplayName("A supported vector still nulls the slot, clearing a value already written") + void supportedVectorClearsExistingValue() throws Exception { + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + IntVector vector = new IntVector("i", allocator)) { + vector.allocateNew(2); + vector.setSafe(0, 42); + vector.setValueCount(1); + assertThat(vector.isNull(0)).isFalse(); + + invokeSetNull(vector, 0); + + assertThat(vector.isNull(0)) + .as("setNull must clear a slot that already holds a value") + .isTrue(); + } + } + + @Test + @DisplayName("MapVector is matched as a map, not captured by the ListVector branch") + void mapVectorIsNotCapturedByListBranch() throws Exception { + // MapVector extends ListVector, so the branch order matters: were the ListVector case + // first, every map would be nulled through the list path. + assertThat(ListVector.class).isAssignableFrom(MapVector.class); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + MapVector vector = MapVector.empty("m", allocator, false)) { + vector.allocateNew(); + + // Reaching the fallback would throw; a clean return proves a map branch was taken. + invokeSetNull(vector, 0); + + assertThat(vector.isNull(0)).isTrue(); + } + } +} From 2c515582f94c656b168e0e918dcbfe457800db39 Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 14:51:52 +0800 Subject: [PATCH 08/11] feat: expose write.data-storage-version, unblocking MAP support Adding MAP to the type converters on its own would have produced a table that fails on first write. Lance gates Map data on file format 2.2+, and every write path in the connector used the SDK default of 2.1. A spike through the production mergeInsert path confirmed the shape of the problem: a Map column is accepted into the schema at CREATE TABLE on 2.1 and survives a reopen intact, but writing rows is rejected by the Rust encoder. Only DDL looks healthy, which is why the gap was not obvious from the capability list. The option has no default. Hardcoding today's 2.1 would pin the connector to it once the SDK default moves forward, so leaving it unset defers to the SDK and keeps the previous behaviour byte for byte. It is applied in LanceCatalog#createTable as well as on both sinks, because the format version is fixed when the dataset is created -- setting it only on the sink leaves an already-materialized table unable to hold a map. That catalog call site turned out to have no test coverage at all: disabling it left every existing suite green. The new ITCase compares an explicitly-2.2 table against an unset one via Dataset#getLanceFileFormatVersion, and disabling the call now fails it. Two things surfaced while testing. mergeInsert does not preserve input order -- the first round-trip assertions read rows by position and saw empty/NULL/populated instead of the written order -- so assertions are keyed by id. And the factory redefines 16 options that LanceOptions already declares, with identical keys; that is the A7 dual-source-of-truth pattern recurring at larger scale. The new option follows the existing pattern in both classes rather than smuggling a refactor into this change. MAP itself is still unmapped. When it lands, flinkTypeToArrowField has to reject a MAP column on a pre-2.2 dataset with a clear message instead of letting the user hit the encoder error. --- ...sue-map-type-blocked-by-storage-version.md | 80 ++++++ docs/src/config.md | 10 + .../flink/connector/lance/LanceSink.java | 10 +- .../connector/lance/LanceUpsertSink.java | 17 +- .../connector/lance/config/LanceOptions.java | 42 ++- .../connector/lance/table/LanceCatalog.java | 12 +- .../lance/table/LanceDynamicTableFactory.java | 11 + .../lance/DataStorageVersionMapITCase.java | 269 ++++++++++++++++++ .../WriteDataStorageVersionOptionTest.java | 87 ++++++ 9 files changed, 532 insertions(+), 6 deletions(-) create mode 100644 .gh-comments/issue-map-type-blocked-by-storage-version.md create mode 100644 src/test/java/org/apache/flink/connector/lance/DataStorageVersionMapITCase.java create mode 100644 src/test/java/org/apache/flink/connector/lance/config/WriteDataStorageVersionOptionTest.java diff --git a/.gh-comments/issue-map-type-blocked-by-storage-version.md b/.gh-comments/issue-map-type-blocked-by-storage-version.md new file mode 100644 index 0000000..ce4b855 --- /dev/null +++ b/.gh-comments/issue-map-type-blocked-by-storage-version.md @@ -0,0 +1,80 @@ +# MAP 类型支持的前置阻塞：Lance 文件格式版本 + +## 结论 + +`MAP` 映射不能只补类型转换层。Lance 的 Map **数据写入**要求文件格式 2.2+，而连接器 +当前所有写入路径都使用 SDK 默认的 2.1，因此照清单补完 `LanceTypeConverter` 与 +`RowDataConverter` 只会得到一个「建表成功、写入必失败」的功能——比不做更糟。 + +## 实测证据 + +spike 走生产 `mergeInsert` 路径（复用 `ArrowArrayStreamsTestAccess`，而非自建 harness）。 + +**默认 2.1**：schema 可建、可往返，写数据时在 Rust 编码层被拒： + +``` +UnsupportedOperationException: Not supported: Map data type is only supported in +Lance file format 2.2+, current version: 2.1 + at lance-encoding/src/encoder.rs:485 +``` + +注意失败点在**写入**而不是建表。`Dataset.create(schema)` 带 Map 列完全成功，重开后 +schema 也无损：`Map(false)>`。 +这正是这个缺口危险的地方——任何只验证 DDL 的测试都看不出问题。 + +**显式 `withDataStorageVersion("2.2")`**：三种场景全部无损往返： + +| 行 | 写入 | 读回 | isNull | +|---|---|---|---| +| 0 | `{"a":1,"b":2}` | `[{"key":"a","value":1},{"key":"b","value":2}]` | false | +| 1 | 空 map | `[]` | false | +| 2 | NULL | `null` | true | + +第 2 行实证了 NULL map 走 `setNull` 路径正确，即上一轮 `MapVector` 分支顺序修正 +（`MapVector extends ListVector`，必须先匹配）所保护的场景。 + +## 当前代码状态 + +三处 `WriteParams` 构造点均未设置存储版本，全部走默认 2.1： + +- `LanceUpsertSink.java:231`（open-or-create） +- `LanceCatalog.java:595`（CREATE TABLE 物化） +- `LanceSink.java:169`（Fragment 写入，只设了 `maxRowsPerFile`） + +`WriteParams.Builder#withDataStorageVersion(String)` 在 lance-core 7.0.0 存在且有效 +（已实测）。 + +## 因此的工作拆分 + +**前置项：暴露存储版本配置** —— 已完成 + +新增 `write.data-storage-version`，接入全部三处 `WriteParams` 构造点。默认值为 +`noDefaultValue()`，不硬编码 `2.1`——否则 SDK 升级默认版本时我们反而把用户钉死在旧版本。 +未设置时行为与改动前完全一致。 + +`LanceCatalog#createTable` 的接入是必需的而非可选：版本在建表时固化，只在 sink 设置对 +已建表无效。该接入点此前**没有任何测试覆盖**——禁用它全部既有套件依然全绿，因此补了 +`DataStorageVersionMapITCase#createTableAppliesConfiguredVersion`，用 +`Dataset#getLanceFileFormatVersion()` 对比「显式 2.2」与「未设置」两张表，并实测禁用后 +该用例失败。 + +**主体项：MAP 类型映射** —— 待做 + +前置项已就位，可以开始。落地时必须在 `flinkTypeToArrowField` 遇到 MAP 而存储版本 < 2.2 +时给出明确报错，而不是让用户在写入时撞 Rust 层的错。 + +## 过程中另外发现的问题 + +**`mergeInsert` 不保证行序。** 首版往返断言按位置取行，实际读回顺序是 +`空 / NULL / {a,b}`，与写入顺序不同。断言已改为按 `id` 关联。这不是缺陷，但任何按位置 +断言的测试都会随机失败。 + +**工厂与 `LanceOptions` 之间有 16 个选项重复定义**，键名完全重合（`path`、`write.mode`、 +`read.*`、`index.*`、`vector.*` 等）。这是 A7「双份真相」在另一处的同类复发，规模更大。 +本次按既有模式在两处都加了新选项以保持一致，未夹带重构——独立处理更安全。 + +## 未验证项 + +- 2.1 与 2.2 数据集能否在同一路径混用（时间旅行读旧版本）。 +- 2.2 是否影响其他类型的编码或既有索引。 +- `MULTISET`（Flink 侧是 `Map`）是否同样受 2.2 限制——推测受同样限制，未实测。 diff --git a/docs/src/config.md b/docs/src/config.md index fc78309..646dd55 100644 --- a/docs/src/config.md +++ b/docs/src/config.md @@ -30,6 +30,16 @@ catalog types are `'lance'` (directory/S3) and `'lance-namespace'` (dir/rest). | `write.batch-size` | ❌ | 1024 | Write batch size | | `write.mode` | ❌ | append | `append` or `overwrite` | | `write.max-rows-per-file` | ❌ | 1000000 | Maximum rows per data file | +| `write.data-storage-version` | ❌ | *(SDK default)* | Lance file format version for written data files, e.g. `2.2`. A `MAP` column requires 2.2+. Fixed when the dataset is created — changing it later does not upgrade an existing dataset. | + +Leaving `write.data-storage-version` unset lets the Lance SDK choose, which is +why it has no default here: pinning today's default would hold the connector +back once the SDK moves forward. + +Some Arrow types are gated on this version. A `MAP` column is accepted into the +schema at `CREATE TABLE` on any version, but writing rows fails inside the Lance +encoder below 2.2 — so set the option at table creation, not after the first +write attempt. ### Vector index 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 b3635f0..62c56fc 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceSink.java @@ -166,9 +166,13 @@ public void flush() throws IOException { } // Build write parameters - WriteParams writeParams = new WriteParams.Builder() - .withMaxRowsPerFile(options.getWriteMaxRowsPerFile()) - .build(); + WriteParams.Builder writeParamsBuilder = new WriteParams.Builder() + .withMaxRowsPerFile(options.getWriteMaxRowsPerFile()); + String storageVersion = options.getWriteDataStorageVersion(); + if (storageVersion != null && !storageVersion.trim().isEmpty()) { + writeParamsBuilder.withDataStorageVersion(storageVersion.trim()); + } + WriteParams writeParams = writeParamsBuilder.build(); // Create Fragment List fragments = Fragment.write() diff --git a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java index dfc945b..3499714 100644 --- a/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java +++ b/src/main/java/org/apache/flink/connector/lance/LanceUpsertSink.java @@ -214,6 +214,21 @@ public void open(Configuration parameters) throws Exception { LOG.info("Lance Upsert Sink opened, primary keys: {}", primaryKeys); } + /** + * Write params for the create-if-absent path, carrying the configured Lance format version. + * + *

    The version is fixed at creation, so it must be supplied here rather than on each write. + * Left unset when the option is absent, which lets the SDK keep its own default. + */ + private WriteParams buildCreateParams() { + WriteParams.Builder builder = new WriteParams.Builder(); + String version = options.getWriteDataStorageVersion(); + if (version != null && !version.trim().isEmpty()) { + builder.withDataStorageVersion(version.trim()); + } + return builder.build(); + } + /** * Return an open {@link Dataset} handle, creating an empty dataset first if it does not yet * exist. This unifies first-write and steady-state paths so that both go through @@ -228,7 +243,7 @@ private Dataset openOrCreate(String datasetPath) throws IOException { openFailure.getMessage(), datasetPath); try { return Dataset.create( - allocator, datasetPath, arrowSchema, new WriteParams.Builder().build()); + allocator, datasetPath, arrowSchema, buildCreateParams()); } catch (Exception createFailure) { // A peer subtask likely won the create race; try open once more. try { diff --git a/src/main/java/org/apache/flink/connector/lance/config/LanceOptions.java b/src/main/java/org/apache/flink/connector/lance/config/LanceOptions.java index db456b9..9acd63e 100644 --- a/src/main/java/org/apache/flink/connector/lance/config/LanceOptions.java +++ b/src/main/java/org/apache/flink/connector/lance/config/LanceOptions.java @@ -139,6 +139,26 @@ public class LanceOptions implements Serializable { .defaultValue(1000000) .withDescription("Maximum rows per data file, default 1000000"); + /** + * Lance file format version used when writing data files. + * + *

    Deliberately has no default: leaving it unset lets the SDK pick, so a future SDK whose + * default moves forward is not held back by a value hardcoded here. Only set it when a + * specific encoding is required. + * + *

    Some Arrow types are gated on the format version -- a MAP column needs 2.2 or newer, and + * on 2.1 the schema is accepted at CREATE TABLE while the first write fails inside the Rust + * encoder. The version is fixed when the dataset is created, so an existing dataset is not + * upgraded by changing this option. + */ + public static final ConfigOption WRITE_DATA_STORAGE_VERSION = ConfigOptions + .key("write.data-storage-version") + .stringType() + .noDefaultValue() + .withDescription("Lance file format version for written data files, e.g. '2.2'. " + + "Unset leaves the choice to the Lance SDK. A MAP column requires 2.2+. " + + "Fixed at dataset creation; changing it does not upgrade an existing dataset."); + // ==================== Vector Index Configuration ==================== /** @@ -397,6 +417,7 @@ public static MetricType fromValue(String value) { private final int writeBatchSize; private final WriteMode writeMode; private final int writeMaxRowsPerFile; + private final String writeDataStorageVersion; private final IndexType indexType; private final String indexColumn; private final int indexNumPartitions; @@ -426,6 +447,7 @@ private LanceOptions(Builder builder) { this.writeBatchSize = builder.writeBatchSize; this.writeMode = builder.writeMode; this.writeMaxRowsPerFile = builder.writeMaxRowsPerFile; + this.writeDataStorageVersion = builder.writeDataStorageVersion; this.indexType = builder.indexType; this.indexColumn = builder.indexColumn; this.indexNumPartitions = builder.indexNumPartitions; @@ -489,6 +511,13 @@ public int getWriteMaxRowsPerFile() { return writeMaxRowsPerFile; } + /** + * Lance file format version for written data files, or {@code null} to let the SDK decide. + */ + public String getWriteDataStorageVersion() { + return writeDataStorageVersion; + } + public IndexType getIndexType() { return indexType; } @@ -609,6 +638,9 @@ public static LanceOptions fromConfiguration(Configuration config) { builder.writeBatchSize(config.get(WRITE_BATCH_SIZE)); builder.writeMode(WriteMode.fromValue(config.get(WRITE_MODE))); builder.writeMaxRowsPerFile(config.get(WRITE_MAX_ROWS_PER_FILE)); + if (config.contains(WRITE_DATA_STORAGE_VERSION)) { + builder.writeDataStorageVersion(config.get(WRITE_DATA_STORAGE_VERSION)); + } // Index configuration builder.indexType(IndexType.fromValue(config.get(INDEX_TYPE))); @@ -661,6 +693,7 @@ public static class Builder { private int writeBatchSize = 1024; private WriteMode writeMode = WriteMode.APPEND; private int writeMaxRowsPerFile = 1000000; + private String writeDataStorageVersion = null; private IndexType indexType = IndexType.IVF_PQ; private String indexColumn; private int indexNumPartitions = 256; @@ -733,6 +766,11 @@ public Builder writeMode(WriteMode writeMode) { return this; } + public Builder writeDataStorageVersion(String writeDataStorageVersion) { + this.writeDataStorageVersion = writeDataStorageVersion; + return this; + } + public Builder writeMaxRowsPerFile(int writeMaxRowsPerFile) { this.writeMaxRowsPerFile = writeMaxRowsPerFile; return this; @@ -902,6 +940,7 @@ public boolean equals(Object o) { Objects.equals(readLimit, that.readLimit) && writeBatchSize == that.writeBatchSize && writeMaxRowsPerFile == that.writeMaxRowsPerFile && + Objects.equals(writeDataStorageVersion, that.writeDataStorageVersion) && indexNumPartitions == that.indexNumPartitions && indexNumBits == that.indexNumBits && indexMaxLevel == that.indexMaxLevel && @@ -928,7 +967,7 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash(path, readBatchSize, readLimit, readColumns, readFilter, writeBatchSize, writeMode, - writeMaxRowsPerFile, indexType, indexColumn, indexNumPartitions, indexNumSubVectors, + writeMaxRowsPerFile, writeDataStorageVersion, indexType, indexColumn, indexNumPartitions, indexNumSubVectors, indexNumBits, indexMaxLevel, indexM, indexEfConstruction, vectorColumn, vectorMetric, vectorNprobes, vectorEf, vectorRefineFactor, defaultDatabase, warehouse, readVersion, readAsOfTimestamp); @@ -947,6 +986,7 @@ public String toString() { ", writeBatchSize=" + writeBatchSize + ", writeMode=" + writeMode + ", writeMaxRowsPerFile=" + writeMaxRowsPerFile + + ", writeDataStorageVersion=" + writeDataStorageVersion + ", indexType=" + indexType + ", indexColumn='" + indexColumn + '\'' + ", indexNumPartitions=" + indexNumPartitions + 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 a92d9ec..edd9937 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 @@ -19,6 +19,7 @@ package org.apache.flink.connector.lance.table; import org.apache.flink.connector.lance.PrimaryKeyPersistence; +import org.apache.flink.connector.lance.config.LanceOptions; import org.apache.flink.connector.lance.converter.LanceTypeConverter; import org.apache.flink.table.api.DataTypes; import org.apache.flink.table.api.Schema; @@ -591,8 +592,17 @@ public void createTable(ObjectPath tablePath, CatalogBaseTable table, boolean ig // 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. + // The format version is fixed when the dataset is created, so it has to be applied here + // and not only on the sink: a table materialized at the SDK default cannot later accept a + // MAP column, which needs 2.2+. + WriteParams.Builder createParams = new WriteParams.Builder(); + String storageVersion = table.getOptions().get( + LanceOptions.WRITE_DATA_STORAGE_VERSION.key()); + if (storageVersion != null && !storageVersion.trim().isEmpty()) { + createParams.withDataStorageVersion(storageVersion.trim()); + } try (Dataset dataset = Dataset.create( - allocator, datasetPath, arrowSchema, new WriteParams.Builder().build())) { + allocator, datasetPath, arrowSchema, createParams.build())) { PrimaryKeyPersistence.persist(dataset, primaryKeys); } catch (Exception e) { throw new CatalogException("Failed to create table: " + tablePath, e); 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 90ec0e9..276d2e7 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 @@ -117,6 +117,13 @@ public class LanceDynamicTableFactory implements DynamicTableSourceFactory, Dyna .defaultValue(1000000) .withDescription("Maximum rows per file"); + public static final ConfigOption WRITE_DATA_STORAGE_VERSION = ConfigOptions + .key("write.data-storage-version") + .stringType() + .noDefaultValue() + .withDescription("Lance file format version for written data files, e.g. '2.2'. " + + "Unset leaves the choice to the Lance SDK. A MAP column requires 2.2+."); + public static final ConfigOption INDEX_TYPE = ConfigOptions .key("index.type") .stringType() @@ -182,6 +189,7 @@ public Set> optionalOptions() { options.add(WRITE_BATCH_SIZE); options.add(WRITE_MODE); options.add(WRITE_MAX_ROWS_PER_FILE); + options.add(WRITE_DATA_STORAGE_VERSION); options.add(INDEX_TYPE); options.add(INDEX_COLUMN); options.add(INDEX_NUM_PARTITIONS); @@ -309,6 +317,9 @@ private LanceOptions buildLanceOptions(ReadableConfig config, MapThis is the evidence behind {@code write.data-storage-version}: a MAP column is accepted into + * the schema at creation on any version, but writing rows only works on 2.2+. Without the option + * there is no way to create a dataset that can hold a map, so adding MAP to the type converters + * alone would produce a table that fails on first write. + * + *

    The version is fixed when the dataset is created, which is why the connector applies the + * option in {@code LanceCatalog#createTable} as well as on the sinks. + */ +class DataStorageVersionMapITCase { + + @TempDir Path tempDir; + + private static Schema mapSchema(RootAllocator allocator) { + Field idField = new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), null); + try (MapVector probe = MapVector.empty("attrs", allocator, false)) { + UnionMapWriter w = probe.getWriter(); + w.allocate(); + w.setPosition(0); + w.startMap(); + w.startEntry(); + w.key().varChar().writeVarChar("k"); + w.value().integer().writeInt(1); + w.endEntry(); + w.endMap(); + w.setValueCount(1); + probe.setValueCount(1); + return new Schema(Arrays.asList(idField, probe.getField())); + } + } + + /** Writes three rows -- populated map, empty map, NULL map -- through the sink's merge path. */ + private static void writeRows(Dataset ds, RootAllocator allocator) throws Exception { + try (VectorSchemaRoot root = VectorSchemaRoot.create(ds.getSchema(), allocator)) { + BigIntVector ids = (BigIntVector) root.getVector("id"); + MapVector attrs = (MapVector) root.getVector("attrs"); + ids.allocateNew(3); + UnionMapWriter writer = attrs.getWriter(); + writer.allocate(); + + ids.setSafe(0, 100L); + writer.setPosition(0); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("a"); + writer.value().integer().writeInt(1); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("b"); + writer.value().integer().writeInt(2); + writer.endEntry(); + writer.endMap(); + + ids.setSafe(1, 200L); + writer.setPosition(1); + writer.startMap(); + writer.endMap(); + + ids.setSafe(2, 300L); + writer.setPosition(2); + writer.writeNull(); + + writer.setValueCount(3); + ids.setValueCount(3); + attrs.setValueCount(3); + root.setRowCount(3); + + MergeInsertParams params = + new MergeInsertParams(Collections.singletonList("id")) + .withMatchedUpdateAll() + .withNotMatched(MergeInsertParams.WhenNotMatched.InsertAll); + ArrowArrayStreamsTestAccess.mergeInsert(ds, params, allocator, root); + } + } + + @Test + @DisplayName("On the SDK default the map schema is accepted but the first write is rejected") + void defaultVersionAcceptsSchemaButRejectsWrite() { + String path = tempDir.resolve("map_default.lance").toString(); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Schema schema = mapSchema(allocator); + + // CREATE TABLE succeeds, which is exactly why this gap is easy to miss: a test that + // only exercises DDL sees nothing wrong. + Dataset.create(allocator, path, schema, new WriteParams.Builder().build()).close(); + + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build())) { + assertThat(ds.getSchema().findField("attrs").getType()) + .isInstanceOf(ArrowType.Map.class); + + assertThatThrownBy(() -> writeRows(ds, allocator)) + .hasMessageContaining("Map data type is only supported in Lance file " + + "format 2.2+"); + } + } catch (Exception e) { + throw new AssertionError("unexpected failure outside the asserted write", e); + } + } + + @Test + @DisplayName("With 2.2 the map column round-trips, including empty and NULL maps") + void version22RoundTripsMapColumn() throws Exception { + String path = tempDir.resolve("map_22.lance").toString(); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Schema schema = mapSchema(allocator); + Dataset.create( + allocator, + path, + schema, + new WriteParams.Builder().withDataStorageVersion("2.2").build()) + .close(); + + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build())) { + writeRows(ds, allocator); + } + + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build())) { + assertThat(ds.countRows()).isEqualTo(3L); + + // mergeInsert does not preserve input order, so rows are matched by id + // rather than by position. + java.util.Map byId = new java.util.HashMap<>(); + java.util.Set nulls = new java.util.HashSet<>(); + try (org.apache.arrow.vector.ipc.ArrowReader reader = ds.newScan().scanBatches()) { + while (reader.loadNextBatch()) { + VectorSchemaRoot out = reader.getVectorSchemaRoot(); + BigIntVector ids = (BigIntVector) out.getVector("id"); + MapVector attrs = (MapVector) out.getVector("attrs"); + for (int i = 0; i < out.getRowCount(); i++) { + long id = ids.get(i); + if (attrs.isNull(i)) { + nulls.add(id); + } else { + byId.put(id, String.valueOf(attrs.getObject(i))); + } + } + } + } + + assertThat(byId.get(100L)) + .as("a populated map must survive the round-trip") + .contains("\"key\":\"a\"") + .contains("\"key\":\"b\""); + + // An empty map must stay distinct from a NULL map. + assertThat(byId).containsKey(200L); + assertThat(byId.get(200L)).isEqualTo("[]"); + assertThat(nulls).doesNotContain(200L); + + // The NULL row exercises setNull's MapVector branch, which has to precede the + // ListVector branch because MapVector extends ListVector. + assertThat(nulls).contains(300L); + } + } + } + + @Test + @DisplayName("CREATE TABLE stamps the configured format version onto the dataset") + void createTableAppliesConfiguredVersion() throws Exception { + // Covers LanceCatalog#createTable specifically. The version is fixed at creation, so if + // the catalog ignored the option the table could never hold a map no matter what the sink + // later requests. Disabling the catalog's withDataStorageVersion call leaves every other + // suite green, which is why this case exists. + org.apache.flink.connector.lance.table.LanceCatalog catalog = + new org.apache.flink.connector.lance.table.LanceCatalog( + "lance", "default", tempDir.toString()); + catalog.open(); + try { + catalog.createDatabase("db", null, false); + + java.util.Map options = new java.util.HashMap<>(); + options.put("write.data-storage-version", "2.2"); + catalog.createTable( + new org.apache.flink.table.catalog.ObjectPath("db", "v22"), + org.apache.flink.table.catalog.CatalogTable.of( + org.apache.flink.table.api.Schema.newBuilder() + .column("id", org.apache.flink.table.api.DataTypes.BIGINT()) + .build(), + "", + Collections.emptyList(), + options), + false); + + // And a table without the option, to prove the difference comes from the option + // rather than from the SDK having changed its default. + catalog.createTable( + new org.apache.flink.table.catalog.ObjectPath("db", "vdefault"), + org.apache.flink.table.catalog.CatalogTable.of( + org.apache.flink.table.api.Schema.newBuilder() + .column("id", org.apache.flink.table.api.DataTypes.BIGINT()) + .build(), + "", + Collections.emptyList(), + Collections.emptyMap()), + false); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + String configured = + formatVersionOf(allocator, tempDir.resolve("db").resolve("v22")); + String fallback = + formatVersionOf(allocator, tempDir.resolve("db").resolve("vdefault")); + + assertThat(configured) + .as("createTable must pass write.data-storage-version to Lance") + .isEqualTo("2.2"); + assertThat(fallback) + .as("an unset option must leave the SDK default in place") + .isNotEqualTo(configured); + } + } finally { + catalog.close(); + } + } + + private static String formatVersionOf(RootAllocator allocator, Path datasetPath) { + try (Dataset ds = + Dataset.open(allocator, datasetPath.toString(), new ReadOptions.Builder().build())) { + return ds.getLanceFileFormatVersion(); + } + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/config/WriteDataStorageVersionOptionTest.java b/src/test/java/org/apache/flink/connector/lance/config/WriteDataStorageVersionOptionTest.java new file mode 100644 index 0000000..efe070a --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/config/WriteDataStorageVersionOptionTest.java @@ -0,0 +1,87 @@ +/* + * 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.config; + +import org.apache.flink.configuration.Configuration; +import org.apache.flink.connector.lance.table.LanceDynamicTableFactory; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Covers the {@code write.data-storage-version} option. + * + *

    The option exists because some Arrow types are gated on the Lance file format version: a MAP + * column needs 2.2 or newer, and on the SDK default of 2.1 the schema is accepted at CREATE TABLE + * while the first write fails inside the Rust encoder. See + * {@code .gh-comments/issue-map-type-blocked-by-storage-version.md}. + */ +class WriteDataStorageVersionOptionTest { + + @Test + @DisplayName("Unset leaves the version null so the SDK keeps its own default") + void unsetLeavesVersionNull() { + // Deliberately not defaulted to "2.1": hardcoding today's SDK default would pin the + // connector to it once the SDK moves forward. + LanceOptions options = + LanceOptions.fromConfiguration( + Configuration.fromMap( + java.util.Collections.singletonMap("path", "/tmp/x.lance"))); + + assertThat(options.getWriteDataStorageVersion()).isNull(); + assertThat(LanceOptions.WRITE_DATA_STORAGE_VERSION.hasDefaultValue()).isFalse(); + } + + @Test + @DisplayName("An explicit version reaches LanceOptions") + void explicitVersionIsCarried() { + java.util.Map map = new java.util.HashMap<>(); + map.put("path", "/tmp/x.lance"); + map.put("write.data-storage-version", "2.2"); + + LanceOptions options = LanceOptions.fromConfiguration(Configuration.fromMap(map)); + + assertThat(options.getWriteDataStorageVersion()).isEqualTo("2.2"); + } + + @Test + @DisplayName("The factory declares the option, so it is a reserved key rather than a user property") + void factoryDeclaresOption() { + // LanceOptionRegistry derives the reserved-key set from the factories' live declarations. + // An undeclared option would be written into the dataset config as a user TBLPROPERTY. + assertThat(new LanceDynamicTableFactory().optionalOptions()) + .extracting(o -> o.key()) + .contains("write.data-storage-version"); + } + + @Test + @DisplayName("The option participates in equals, hashCode and toString") + void optionParticipatesInValueSemantics() { + // A field missing from equals makes two differently-configured option objects compare + // equal, which would silently reuse a cached sink or catalog built for another version. + LanceOptions base = LanceOptions.builder().path("/tmp/x.lance").build(); + LanceOptions withVersion = + LanceOptions.builder().path("/tmp/x.lance").writeDataStorageVersion("2.2").build(); + + assertThat(withVersion).isNotEqualTo(base); + assertThat(withVersion.hashCode()).isNotEqualTo(base.hashCode()); + assertThat(withVersion.toString()).contains("2.2"); + } +} From 9581d2bbf0e2cf10f30ad85a6ebfd16ce4d885b3 Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 15:42:51 +0800 Subject: [PATCH 09/11] feat: support MAP columns end to end Builds on the write.data-storage-version groundwork. All three conversion directions in LanceTypeConverter and all four dispatch points in RowDataConverter now handle MAP, so a map column can be created, written and read rather than only appearing in a schema. Three things needed verifying rather than assuming, and each was confirmed by disabling the code and watching a test fail: setIndexDefined on the entries struct is required. Removing it leaves the in-memory converter round-trip fully green, because that test reads back the vectors it just wrote and never passes through Lance's validation. A real dataset rejects it immediately: "The field `entries` contained null values even though the field is marked non-null in the schema". The in-memory layer has no discriminating power for non-null constraints, so both test layers are load bearing. The MapVector branch has to precede ListVector on the write path too, not just in setNull. MapVector extends ListVector, so without it a map falls into the list branch. The same ordering applies in reverse: ArrowType.Map must be checked before ArrowType.List, or a map degrades into ARRAY> and stops round-tripping. Key and value types are far narrower than a top-level column: only INT, BIGINT, FLOAT, DOUBLE and STRING, bounded by RowDataConverter's array element helpers. Arrow is perfectly happy with Map, so an unchecked MAP would create a table whose first write fails -- the same trap the 2.2 format requirement set for MAP itself. mapEntriesField rejects it at DDL time instead. LanceCatalog#createTable refuses a MAP column when the configured version is demonstrably below 2.2, naming the column and the option. An unset version is only warned about: the effective default belongs to the SDK, and refusing here would misfire the moment that default reaches 2.2. Arrow does not allow a nullable map key, which makes the obvious DataTypes.MAP(STRING(), INT()) invalid; the message says so explicitly rather than leaving the user to discover .notNull(). Empty and NULL maps are stored distinctly, and a NULL value is allowed where a NULL key is not. The unsupported-type sample in LanceNamespaceCatalogSchemaTest moves for the third time -- DECIMAL, then MAP, now MULTISET -- since each in turn gained a mapping. Its comment records the trail so the next move is mechanical. --- ...sue-map-type-blocked-by-storage-version.md | 39 +++- docs/src/config.md | 26 +++ .../lance/converter/LanceTypeConverter.java | 115 +++++++++++- .../lance/converter/RowDataConverter.java | 116 ++++++++++++ .../connector/lance/table/LanceCatalog.java | 89 +++++++++ .../lance/DataStorageVersionMapITCase.java | 86 +++++++++ .../lance/MapColumnLanceRoundTripITCase.java | 159 ++++++++++++++++ .../converter/LanceTypeConverterMapTest.java | 174 ++++++++++++++++++ .../converter/RowDataConverterMapTest.java | 171 +++++++++++++++++ .../LanceNamespaceCatalogSchemaTest.java | 42 ++++- 10 files changed, 1007 insertions(+), 10 deletions(-) create mode 100644 src/test/java/org/apache/flink/connector/lance/MapColumnLanceRoundTripITCase.java create mode 100644 src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMapTest.java create mode 100644 src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterMapTest.java diff --git a/.gh-comments/issue-map-type-blocked-by-storage-version.md b/.gh-comments/issue-map-type-blocked-by-storage-version.md index ce4b855..aef049f 100644 --- a/.gh-comments/issue-map-type-blocked-by-storage-version.md +++ b/.gh-comments/issue-map-type-blocked-by-storage-version.md @@ -58,10 +58,37 @@ schema 也无损：`Map(false)>` 且不再往返。 + +**key/value 的类型范围远窄于顶层列。** 只支持 INT/BIGINT/FLOAT/DOUBLE/STRING，受 +`RowDataConverter` 的 `readArrayData`/`writeArrayData` 限制（ARRAY 元素同此约束）。 +`MAP` 在 Arrow 层完全合法、建表会放行，但首次写入必失败——与 MAP 本身踩的 +是同一个坑，因此在 `mapEntriesField` 里提前拒绝。要放宽得先扩那两个 helper。 + +## 用户侧易踩点 + +`DataTypes.MAP(STRING(), INT())` 的 key 默认可空，而 Arrow 不允许可空 key，所以最自然的 +写法会被拒绝。正确写法是 `DataTypes.MAP(DataTypes.STRING().notNull(), DataTypes.INT())`。 +错误信息显式说明了这一点。 ## 过程中另外发现的问题 @@ -73,8 +100,12 @@ schema 也无损：`Map(false)`）是否同样受 2.2 限制——推测受同样限制，未实测。 +- `MULTISET`（Flink 侧是 `Map`）仍无映射，推测同受 2.2 限制，未实测。 diff --git a/docs/src/config.md b/docs/src/config.md index 646dd55..7f2ae18 100644 --- a/docs/src/config.md +++ b/docs/src/config.md @@ -41,6 +41,32 @@ schema at `CREATE TABLE` on any version, but writing rows fails inside the Lance encoder below 2.2 — so set the option at table creation, not after the first write attempt. +### MAP columns + +A `MAP` column requires `write.data-storage-version` to be `2.2` or newer. +`CREATE TABLE` fails fast if the option is set to something older, naming the +column. If the option is unset the connector only logs a warning, since the +effective default belongs to the Lance SDK. + +Two constraints apply to the key and value types: + +```sql +-- Rejected: DataTypes.MAP(STRING(), INT()) yields a nullable key, and Arrow +-- does not allow one. +CREATE TABLE t (attrs MAP) WITH (...); + +-- Correct: the key is declared NOT NULL. +CREATE TABLE t (attrs MAP) WITH ( + 'write.data-storage-version' = '2.2', ... +); +``` + +Keys and values support `INT`, `BIGINT`, `FLOAT`, `DOUBLE` and `STRING`. This is +narrower than what a top-level column accepts — the same limit applies to +`ARRAY` elements — and a type outside it is rejected at `CREATE TABLE` rather +than on the first write. An empty map and a `NULL` map are stored distinctly. A +`NULL` value is allowed; a `NULL` key is not. + ### Vector index | Option | Required | Default | Description | diff --git a/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java b/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java index a0a1aff..089bd13 100644 --- a/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java +++ b/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java @@ -31,6 +31,7 @@ import org.apache.flink.table.types.logical.IntType; import org.apache.flink.table.types.logical.LocalZonedTimestampType; import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.table.types.logical.SmallIntType; import org.apache.flink.table.types.logical.TimeType; @@ -42,6 +43,7 @@ import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; @@ -193,6 +195,26 @@ public static LogicalType arrowTypeToFlinkType(Field field) { return new ArrayType(nullable, elementType); } throw new UnsupportedTypeException("FixedSizeList must contain child type"); + } else if (arrowType instanceof ArrowType.Map) { + // Must precede the List branch. An Arrow map is physically a list of entry structs, so + // a List-first check would map it to ARRAY> and the column would stop + // round-tripping as a MAP. + List children = field.getChildren(); + if (children == null || children.isEmpty()) { + throw new UnsupportedTypeException("Map must contain an entries child"); + } + Field entries = children.get(0); + List keyValue = entries.getChildren(); + if (keyValue == null || keyValue.size() != 2) { + throw new UnsupportedTypeException( + "Map entries must contain exactly key and value children, found " + + (keyValue == null ? 0 : keyValue.size())); + } + LogicalType keyType = arrowTypeToFlinkType(keyValue.get(0)); + LogicalType valueType = arrowTypeToFlinkType(keyValue.get(1)); + // Arrow guarantees a non-null key; carry that through so a round-trip does not hand + // back a MAP the converter would then refuse on the way in. + return new MapType(nullable, keyType.copy(false), valueType); } else if (arrowType instanceof ArrowType.List || arrowType instanceof ArrowType.LargeList) { // Regular list type List children = field.getChildren(); @@ -282,6 +304,22 @@ public static Field flinkTypeToArrowField(String name, LogicalType logicalType) children.add(childField); // For vector types, use List type arrowType = ArrowType.List.INSTANCE; + } else if (logicalType instanceof MapType) { + MapType mapType = (MapType) logicalType; + LogicalType keyType = mapType.getKeyType(); + // Arrow requires map keys to be non-null, while Flink's MapType allows a nullable key + // type. Silently widening it would let a NULL key reach the encoder, so reject it here + // where the message can still name the offending column. + if (keyType.isNullable()) { + throw new UnsupportedTypeException( + "MAP key must be NOT NULL for column '" + name + "': Arrow map keys cannot " + + "be nullable. Declare the key as e.g. MAP."); + } + children = new ArrayList<>(); + children.add(mapEntriesField(name, keyType, mapType.getValueType())); + // keysSorted=false: nothing in the write path sorts entries, and claiming otherwise + // would let a reader skip its own ordering work on unordered data. + arrowType = new ArrowType.Map(false); } else if (logicalType instanceof RowType) { RowType rowType = (RowType) logicalType; children = new ArrayList<>(); @@ -299,7 +337,76 @@ public static Field flinkTypeToArrowField(String name, LogicalType logicalType) } /** - * Create vector field (FixedSizeList) + * Build the {@code entries} struct that backs an Arrow map field. + * + *

    Arrow fixes both the names and the nullability here: the struct is called {@code entries} + * and must itself be non-nullable, and {@code key} must be non-nullable. Only {@code value} may + * be null. Getting any of that wrong surfaces later as an opaque IPC or encoder error, so the + * names come from {@link MapVector}'s constants rather than string literals -- note that + * {@code MapVector.DATA_VECTOR_NAME} is {@code entries}, whereas the inherited + * {@code BaseRepeatedValueVector.DATA_VECTOR_NAME} is {@code $data$}. + */ + private static Field mapEntriesField( + String mapColumnName, LogicalType keyType, LogicalType valueType) { + // Arrow accepts far more element types here than the read/write path can actually move, so + // an unchecked MAP would create a table whose first write fails deep in the + // converter. Reject it at DDL time instead, where the message can name the column. + requireSupportedMapElement(mapColumnName, "key", keyType); + requireSupportedMapElement(mapColumnName, "value", valueType); + + Field keyField = flinkTypeToArrowField(MapVector.KEY_NAME, keyType); + if (keyField.isNullable()) { + // Defensive: the caller already rejects a nullable key type, but a converter that + // widened nullability on the way out would otherwise produce a schema Arrow refuses. + keyField = + new Field( + MapVector.KEY_NAME, + new FieldType(false, keyField.getType(), null), + keyField.getChildren()); + } + Field valueField = flinkTypeToArrowField(MapVector.VALUE_NAME, valueType); + + List entryChildren = new ArrayList<>(); + entryChildren.add(keyField); + entryChildren.add(valueField); + + return new Field( + MapVector.DATA_VECTOR_NAME, + new FieldType(false, ArrowType.Struct.INSTANCE, null), + entryChildren); + } + + /** + * Element types the map read/write path can carry. + * + *

    This mirrors what {@code RowDataConverter}'s array element helpers implement, since map + * keys and values reuse them. It is narrower than the set of types allowed for a top-level + * column, and widening it means extending those helpers first. + */ + private static void requireSupportedMapElement( + String mapColumnName, String role, LogicalType elementType) { + boolean supported = + elementType instanceof IntType + || elementType instanceof BigIntType + || elementType instanceof FloatType + || elementType instanceof DoubleType + || elementType instanceof VarCharType; + if (!supported) { + throw new UnsupportedTypeException( + "Unsupported MAP " + + role + + " type for column '" + + mapColumnName + + "': " + + elementType.getClass().getSimpleName() + + ". MAP keys and values support INT, BIGINT, FLOAT, DOUBLE and STRING. " + + "Arrow would accept more, but the connector's map read/write path " + + "would then fail on the first write rather than here."); + } + } + + /** + * Create vector field (FixedSizeList<Float32>) * * @param name Field name * @param dimension Vector dimension @@ -431,6 +538,12 @@ public static DataType toDataType(LogicalType logicalType) { ArrayType arrayType = (ArrayType) logicalType; DataType elementDataType = toDataType(arrayType.getElementType()); return DataTypes.ARRAY(elementDataType); + } else if (logicalType instanceof MapType) { + MapType mapType = (MapType) logicalType; + DataType keyDataType = toDataType(mapType.getKeyType()); + DataType valueDataType = toDataType(mapType.getValueType()); + // The key stays NOT NULL to match Arrow, which does not allow a nullable map key. + return DataTypes.MAP(keyDataType.notNull(), valueDataType); } else if (logicalType instanceof RowType) { RowType rowType = (RowType) logicalType; DataTypes.Field[] fields = rowType.getFields().stream() diff --git a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java index a102bd4..5059c8c 100644 --- a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java +++ b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java @@ -20,6 +20,8 @@ import org.apache.flink.table.data.ArrayData; import org.apache.flink.table.data.DecimalData; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.MapData; import org.apache.flink.table.data.GenericArrayData; import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; @@ -36,6 +38,7 @@ import org.apache.flink.table.types.logical.IntType; import org.apache.flink.table.types.logical.LocalZonedTimestampType; import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.table.types.logical.SmallIntType; import org.apache.flink.table.types.logical.TimeType; @@ -85,6 +88,8 @@ import java.time.Instant; import java.time.LocalDate; import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.Map; import java.util.List; /** @@ -225,6 +230,8 @@ private Object readValue(FieldVector vector, int index, LogicalType logicalType) value, decimalType.getPrecision(), decimalType.getScale()); } else if (logicalType instanceof ArrayType) { return readArray(vector, index, (ArrayType) logicalType); + } else if (logicalType instanceof MapType) { + return readMap(vector, index, (MapType) logicalType); } else if (logicalType instanceof RowType) { return readStruct(vector, index, (RowType) logicalType); } @@ -332,6 +339,64 @@ private ArrayData readArray(FieldVector vector, int index, ArrayType arrayType) "Unsupported array Vector type: " + vector.getClass().getSimpleName()); } + /** + * Read map value. + * + *

    An Arrow map is a list of non-nullable {@code entries} structs, each holding a {@code key} + * and a {@code value} child. The offsets come from the enclosing list, so the key and value + * slices are read from the same index range. + */ + private MapData readMap(FieldVector vector, int index, MapType mapType) { + if (!(vector instanceof MapVector)) { + // Note this cannot be relaxed to ListVector: a plain list carries no key child, and + // treating one as a map would read garbage out of the element vector. + throw new LanceTypeConverter.UnsupportedTypeException( + "Unsupported map Vector type: " + vector.getClass().getSimpleName()); + } + + MapVector mapVector = (MapVector) vector; + int startIndex = mapVector.getElementStartIndex(index); + int endIndex = mapVector.getElementEndIndex(index); + int size = endIndex - startIndex; + + StructVector entries = (StructVector) mapVector.getDataVector(); + FieldVector keyVector = entries.getChild(MapVector.KEY_NAME); + FieldVector valueVector = entries.getChild(MapVector.VALUE_NAME); + if (keyVector == null || valueVector == null) { + throw new LanceTypeConverter.UnsupportedTypeException( + "Map entries struct must expose '" + MapVector.KEY_NAME + "' and '" + + MapVector.VALUE_NAME + "' children"); + } + + ArrayData keys = readArrayData(keyVector, startIndex, size, mapType.getKeyType()); + ArrayData values = readArrayData(valueVector, startIndex, size, mapType.getValueType()); + return new GenericMapData(toJavaMap(keys, values, mapType)); + } + + /** + * Materialize key/value slices into the map {@link GenericMapData} expects. + * + *

    Duplicate keys collapse to the last occurrence, matching how Flink's own map + * implementations behave when a duplicate reaches them. + */ + private Map toJavaMap(ArrayData keys, ArrayData values, MapType mapType) { + Map result = new LinkedHashMap<>(); + for (int i = 0; i < keys.size(); i++) { + Object key = elementAt(keys, i, mapType.getKeyType()); + Object value = elementAt(values, i, mapType.getValueType()); + result.put(key, value); + } + return result; + } + + /** Extract one element out of an {@link ArrayData} as the object GenericMapData stores. */ + private Object elementAt(ArrayData array, int i, LogicalType elementType) { + if (array.isNullAt(i)) { + return null; + } + return ArrayData.createElementGetter(elementType).getElementOrNull(array, i); + } + /** * Read array data */ @@ -460,6 +525,8 @@ private Object getFieldValue(RowData rowData, int index, LogicalType logicalType return rowData.getDecimal(index, decimalType.getPrecision(), decimalType.getScale()); } else if (logicalType instanceof ArrayType) { return rowData.getArray(index); + } else if (logicalType instanceof MapType) { + return rowData.getMap(index); } else if (logicalType instanceof RowType) { RowType nestedRowType = (RowType) logicalType; return rowData.getRow(index, nestedRowType.getFieldCount()); @@ -511,6 +578,8 @@ private void writeValue(FieldVector vector, int index, Object value, LogicalType ((DecimalVector) vector).setSafe(index, ((DecimalData) value).toBigDecimal()); } else if (logicalType instanceof ArrayType) { writeArray(vector, index, (ArrayData) value, (ArrayType) logicalType); + } else if (logicalType instanceof MapType) { + writeMap(vector, index, (MapData) value, (MapType) logicalType); } else if (logicalType instanceof RowType) { writeStruct(vector, index, (RowData) value, (RowType) logicalType); } else { @@ -699,6 +768,53 @@ private void writeArray(FieldVector vector, int index, ArrayData arrayData, Arra } } + /** + * Write map value. + * + *

    Mirrors {@link #writeArray}'s list handling for the offsets, with two additions specific + * to maps: each {@code entries} slot has to be marked defined or the struct reads back as NULL + * even though key and value were written, and a NULL key is rejected because Arrow does not + * allow one. + */ + private void writeMap(FieldVector vector, int index, MapData mapData, MapType mapType) { + if (!(vector instanceof MapVector)) { + throw new LanceTypeConverter.UnsupportedTypeException( + "Unsupported map Vector type: " + vector.getClass().getSimpleName()); + } + + MapVector mapVector = (MapVector) vector; + ArrayData keys = mapData.keyArray(); + ArrayData values = mapData.valueArray(); + int size = mapData.size(); + + for (int i = 0; i < size; i++) { + if (keys.isNullAt(i)) { + throw new IllegalArgumentException( + "MAP key must not be NULL: Arrow map keys are non-nullable, so a NULL key " + + "cannot be written (entry " + i + ")"); + } + } + + mapVector.startNewValue(index); + + StructVector entries = (StructVector) mapVector.getDataVector(); + FieldVector keyVector = entries.getChild(MapVector.KEY_NAME); + FieldVector valueVector = entries.getChild(MapVector.VALUE_NAME); + int startIndex = mapVector.getElementStartIndex(index); + + writeArrayData(keyVector, startIndex, keys, mapType.getKeyType()); + writeArrayData(valueVector, startIndex, values, mapType.getValueType()); + + // Without this the entries struct keeps a zero validity bit and the whole entry reads back + // as NULL, which looks like data loss rather than a missing flag. + for (int i = 0; i < size; i++) { + entries.setIndexDefined(startIndex + i); + } + entries.setValueCount(startIndex + size); + + mapVector.endValue(index, size); + } + /** * Write array data */ 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 edd9937..2e757ff 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 @@ -601,6 +601,7 @@ public void createTable(ObjectPath tablePath, CatalogBaseTable table, boolean ig if (storageVersion != null && !storageVersion.trim().isEmpty()) { createParams.withDataStorageVersion(storageVersion.trim()); } + rejectMapColumnsOnUnsupportedVersion(arrowSchema, storageVersion, tablePath); try (Dataset dataset = Dataset.create( allocator, datasetPath, arrowSchema, createParams.build())) { PrimaryKeyPersistence.persist(dataset, primaryKeys); @@ -616,6 +617,94 @@ public void createTable(ObjectPath tablePath, CatalogBaseTable table, boolean ig LOG.info("Created table: {} (primary keys: {})", tablePath, primaryKeys); } + /** + * Fail early when a MAP column is created on a dataset format that cannot hold one. + * + *

    Lance accepts a map into the schema on any version and only rejects it when the first row + * is written, deep in the Rust encoder. Catching it here means the error names the table and + * the option to set. + * + *

    An unset version is only warned about, not rejected: the effective default belongs to the + * SDK, so refusing here would break the moment that default moves to 2.2 or later. + */ + private void rejectMapColumnsOnUnsupportedVersion( + org.apache.arrow.vector.types.pojo.Schema arrowSchema, + String storageVersion, + ObjectPath tablePath) { + List mapColumns = new ArrayList<>(); + for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) { + if (containsMap(field)) { + mapColumns.add(field.getName()); + } + } + if (mapColumns.isEmpty()) { + return; + } + + String configured = storageVersion == null ? null : storageVersion.trim(); + if (configured == null || configured.isEmpty()) { + LOG.warn( + "Table {} declares MAP column(s) {} without {}. MAP data requires Lance format " + + "2.2+; if the SDK default is older the first write will fail in the " + + "encoder.", + tablePath.getFullName(), + mapColumns, + LanceOptions.WRITE_DATA_STORAGE_VERSION.key()); + return; + } + + if (isBelow22(configured)) { + throw new CatalogException( + "Table " + + tablePath.getFullName() + + " declares MAP column(s) " + + mapColumns + + " but " + + LanceOptions.WRITE_DATA_STORAGE_VERSION.key() + + " is '" + + configured + + "'. MAP data requires Lance format 2.2 or newer; the schema would be " + + "accepted here and the first write would then fail in the encoder."); + } + } + + /** Whether a field is a map, or transitively contains one. */ + private static boolean containsMap(org.apache.arrow.vector.types.pojo.Field field) { + if (field.getType() instanceof org.apache.arrow.vector.types.pojo.ArrowType.Map) { + return true; + } + List children = field.getChildren(); + if (children == null) { + return false; + } + for (org.apache.arrow.vector.types.pojo.Field child : children) { + if (containsMap(child)) { + return true; + } + } + return false; + } + + /** + * Compare a {@code major.minor} version string against 2.2. + * + *

    Anything unparseable is treated as new enough, so an alias such as {@code stable} is left + * for Lance to accept or reject rather than guessed at here. + */ + private static boolean isBelow22(String version) { + String[] parts = version.split("\\."); + if (parts.length < 2) { + return false; + } + try { + int major = Integer.parseInt(parts[0].trim()); + int minor = Integer.parseInt(parts[1].trim()); + return major < 2 || (major == 2 && minor < 2); + } catch (NumberFormatException e) { + return false; + } + } + /** * Primary {@code ALTER TABLE} entry point: applies the planner's explicit {@link TableChange} * list instead of reconstructing intent from a schema diff. diff --git a/src/test/java/org/apache/flink/connector/lance/DataStorageVersionMapITCase.java b/src/test/java/org/apache/flink/connector/lance/DataStorageVersionMapITCase.java index 712f7f1..2a0ace1 100644 --- a/src/test/java/org/apache/flink/connector/lance/DataStorageVersionMapITCase.java +++ b/src/test/java/org/apache/flink/connector/lance/DataStorageVersionMapITCase.java @@ -266,4 +266,90 @@ private static String formatVersionOf(RootAllocator allocator, Path datasetPath) return ds.getLanceFileFormatVersion(); } } + + @Test + @DisplayName("CREATE TABLE with a MAP column on a pre-2.2 version fails before the table exists") + void mapColumnOnOldVersionIsRejectedAtDdlTime() throws Exception { + // Without this check the schema is accepted, the table is materialized, and the user only + // finds out on the first write via a Rust-level encoder message that names neither the + // table nor the option to set. + org.apache.flink.connector.lance.table.LanceCatalog catalog = + new org.apache.flink.connector.lance.table.LanceCatalog( + "lance", "default", tempDir.toString()); + catalog.open(); + try { + catalog.createDatabase("db", null, false); + + java.util.Map options = new java.util.HashMap<>(); + options.put("write.data-storage-version", "2.1"); + + org.apache.flink.table.catalog.ObjectPath tablePath = + new org.apache.flink.table.catalog.ObjectPath("db", "bad_map"); + org.apache.flink.table.catalog.CatalogTable table = + org.apache.flink.table.catalog.CatalogTable.of( + org.apache.flink.table.api.Schema.newBuilder() + .column("id", org.apache.flink.table.api.DataTypes.INT()) + .column( + "attrs", + org.apache.flink.table.api.DataTypes.MAP( + org.apache.flink.table.api.DataTypes.STRING() + .notNull(), + org.apache.flink.table.api.DataTypes.INT())) + .build(), + "", + Collections.emptyList(), + options); + + org.assertj.core.api.Assertions.assertThatThrownBy( + () -> catalog.createTable(tablePath, table, false)) + .hasMessageContaining("attrs") + .hasMessageContaining("2.2"); + + // The failure must come before materialization, otherwise a half-created table is left + // behind for the next CREATE to trip over. + assertThat(catalog.tableExists(tablePath)) + .as("the table must not be left behind after a rejected CREATE") + .isFalse(); + } finally { + catalog.close(); + } + } + + @Test + @DisplayName("CREATE TABLE with a MAP column on 2.2 is accepted") + void mapColumnOn22IsAccepted() throws Exception { + org.apache.flink.connector.lance.table.LanceCatalog catalog = + new org.apache.flink.connector.lance.table.LanceCatalog( + "lance", "default", tempDir.toString()); + catalog.open(); + try { + catalog.createDatabase("db", null, false); + + java.util.Map options = new java.util.HashMap<>(); + options.put("write.data-storage-version", "2.2"); + + org.apache.flink.table.catalog.ObjectPath tablePath = + new org.apache.flink.table.catalog.ObjectPath("db", "good_map"); + catalog.createTable( + tablePath, + org.apache.flink.table.catalog.CatalogTable.of( + org.apache.flink.table.api.Schema.newBuilder() + .column("id", org.apache.flink.table.api.DataTypes.INT()) + .column( + "attrs", + org.apache.flink.table.api.DataTypes.MAP( + org.apache.flink.table.api.DataTypes.STRING() + .notNull(), + org.apache.flink.table.api.DataTypes.INT())) + .build(), + "", + Collections.emptyList(), + options), + false); + + assertThat(catalog.tableExists(tablePath)).isTrue(); + } finally { + catalog.close(); + } + } } diff --git a/src/test/java/org/apache/flink/connector/lance/MapColumnLanceRoundTripITCase.java b/src/test/java/org/apache/flink/connector/lance/MapColumnLanceRoundTripITCase.java new file mode 100644 index 0000000..25beec7 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/MapColumnLanceRoundTripITCase.java @@ -0,0 +1,159 @@ +/* + * 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.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.flink.connector.lance.converter.LanceTypeConverter; +import org.apache.flink.connector.lance.converter.RowDataConverter; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.MapData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.lance.Dataset; +import org.lance.ReadOptions; +import org.lance.WriteParams; +import org.lance.merge.MergeInsertParams; + +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * End-to-end MAP support: Flink rows through {@link RowDataConverter} into a real Lance dataset and + * back. + * + *

    The in-memory converter test cannot prove much on its own, because it reads back from the same + * vectors it wrote. Only a real dataset exercises Lance's own encode and decode, which is where the + * 2.2 format requirement and the entries-struct validity actually bite. + */ +class MapColumnLanceRoundTripITCase { + + @TempDir Path tempDir; + + private static final RowType ROW_TYPE = + RowType.of( + new LogicalType[] { + new IntType(false), + new MapType( + true, + new VarCharType(false, VarCharType.MAX_LENGTH), + new IntType(true)) + }, + new String[] {"id", "attrs"}); + + private static MapData map(Object... kv) { + Map m = new HashMap<>(); + for (int i = 0; i < kv.length; i += 2) { + m.put(StringData.fromString((String) kv[i]), kv[i + 1]); + } + return new GenericMapData(m); + } + + @Test + @DisplayName("MAP rows written through the converter survive a real Lance round-trip") + void mapSurvivesLanceRoundTrip() throws Exception { + String path = tempDir.resolve("map_rows").toString(); + + List input = + Arrays.asList( + GenericRowData.of(1, map("a", 10, "b", 20)), + GenericRowData.of(2, map()), + GenericRowData.of(3, null), + GenericRowData.of(4, map("only", 7))); + + RowDataConverter converter = new RowDataConverter(ROW_TYPE); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + // MAP data needs Lance format 2.2+; on the SDK default the write fails in the encoder. + Dataset.create( + allocator, + path, + LanceTypeConverter.toArrowSchema(ROW_TYPE), + new WriteParams.Builder().withDataStorageVersion("2.2").build()) + .close(); + + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build()); + VectorSchemaRoot root = converter.createVectorSchemaRoot(allocator)) { + converter.toVectorSchemaRoot(input, root); + MergeInsertParams params = + new MergeInsertParams(Collections.singletonList("id")) + .withMatchedUpdateAll() + .withNotMatched(MergeInsertParams.WhenNotMatched.InsertAll); + ArrowArrayStreamsTestAccess.mergeInsert(ds, params, allocator, root); + } + + // Reopen so the assertions read committed data rather than the pre-merge snapshot. + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build())) { + assertThat(ds.countRows()).isEqualTo(4L); + + List readBack = new ArrayList<>(); + try (org.apache.arrow.vector.ipc.ArrowReader reader = ds.newScan().scanBatches()) { + while (reader.loadNextBatch()) { + readBack.addAll(converter.toRowDataList(reader.getVectorSchemaRoot())); + } + } + + // mergeInsert does not preserve input order, so rows are keyed by id. + Map byId = new HashMap<>(); + for (RowData row : readBack) { + byId.put(row.getInt(0), row); + } + assertThat(byId).hasSize(4); + + Map first = flatten(byId.get(1).getMap(1)); + assertThat(first).containsEntry("a", 10).containsEntry("b", 20); + + // An empty map must stay distinct from a NULL map through Lance's own encoding. + assertThat(byId.get(2).isNullAt(1)).isFalse(); + assertThat(byId.get(2).getMap(1).size()).isZero(); + + assertThat(byId.get(3).isNullAt(1)).isTrue(); + + assertThat(flatten(byId.get(4).getMap(1))).containsExactlyEntriesOf( + Collections.singletonMap("only", 7)); + } + } + } + + private static Map flatten(MapData mapData) { + Map out = new HashMap<>(); + for (int i = 0; i < mapData.size(); i++) { + out.put( + mapData.keyArray().getString(i).toString(), + mapData.valueArray().isNullAt(i) ? null : mapData.valueArray().getInt(i)); + } + return out; + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMapTest.java b/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMapTest.java new file mode 100644 index 0000000..a369348 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMapTest.java @@ -0,0 +1,174 @@ +/* + * 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.converter; + +import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.flink.table.types.logical.ArrayType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Covers the MAP mapping in {@link LanceTypeConverter}. */ +class LanceTypeConverterMapTest { + + private static MapType mapOf(LogicalType key, LogicalType value) { + return new MapType(true, key, value); + } + + private static LogicalType notNullString() { + return new VarCharType(false, VarCharType.MAX_LENGTH); + } + + @Test + @DisplayName("A MAP column becomes an Arrow map with the entries/key/value shape Arrow requires") + void mapBecomesArrowMapWithEntriesStruct() { + Field field = + LanceTypeConverter.flinkTypeToArrowField( + "attrs", mapOf(notNullString(), new IntType(true))); + + assertThat(field.getType()).isInstanceOf(ArrowType.Map.class); + assertThat(field.getChildren()).hasSize(1); + + Field entries = field.getChildren().get(0); + // The names are fixed by Arrow. MapVector.DATA_VECTOR_NAME is "entries", while the + // inherited BaseRepeatedValueVector constant is "$data$" -- using the wrong one produces a + // schema that fails much later. + assertThat(entries.getName()).isEqualTo("entries"); + assertThat(entries.getType()).isInstanceOf(ArrowType.Struct.class); + assertThat(entries.isNullable()) + .as("the entries struct itself must be non-nullable") + .isFalse(); + + assertThat(entries.getChildren()).hasSize(2); + Field key = entries.getChildren().get(0); + Field value = entries.getChildren().get(1); + assertThat(key.getName()).isEqualTo(MapVector.KEY_NAME); + assertThat(value.getName()).isEqualTo(MapVector.VALUE_NAME); + assertThat(key.isNullable()).as("Arrow map keys cannot be nullable").isFalse(); + assertThat(value.isNullable()).as("map values may be null").isTrue(); + } + + @Test + @DisplayName("A nullable MAP key is rejected up front, naming the column") + void nullableKeyIsRejected() { + // Arrow forbids a nullable key. Widening it silently would push the failure down into the + // encoder, where the message no longer says which column is at fault. + assertThatThrownBy( + () -> + LanceTypeConverter.flinkTypeToArrowField( + "attrs", + mapOf( + new VarCharType(true, VarCharType.MAX_LENGTH), + new IntType(true)))) + .isInstanceOf(LanceTypeConverter.UnsupportedTypeException.class) + .hasMessageContaining("attrs") + .hasMessageContaining("NOT NULL"); + } + + @Test + @DisplayName("An Arrow map converts back to MAP rather than ARRAY>") + void arrowMapConvertsBackToMap() { + // An Arrow map is physically a list of entry structs, so a List-first check in + // arrowTypeToFlinkType would silently degrade the column to ARRAY> and the schema + // would stop round-tripping. + Field field = + LanceTypeConverter.flinkTypeToArrowField( + "attrs", mapOf(notNullString(), new IntType(true))); + + LogicalType back = LanceTypeConverter.arrowTypeToFlinkType(field); + + assertThat(back).isInstanceOf(MapType.class); + assertThat(back).isNotInstanceOf(ArrayType.class); + + MapType mapType = (MapType) back; + assertThat(mapType.getKeyType()).isInstanceOf(VarCharType.class); + assertThat(mapType.getValueType()).isInstanceOf(IntType.class); + assertThat(mapType.getKeyType().isNullable()) + .as("the non-null key must survive the round-trip, or the reverse conversion " + + "would produce a MAP the forward conversion then refuses") + .isFalse(); + } + + @Test + @DisplayName("A MAP nested inside a ROW round-trips") + void mapNestedInRowRoundTrips() { + RowType rowType = + RowType.of( + new LogicalType[] {new IntType(true), mapOf(notNullString(), new IntType(true))}, + new String[] {"id", "attrs"}); + + Field field = LanceTypeConverter.flinkTypeToArrowField("payload", rowType); + LogicalType back = LanceTypeConverter.arrowTypeToFlinkType(field); + + assertThat(back).isInstanceOf(RowType.class); + LogicalType nested = ((RowType) back).getTypeAt(1); + assertThat(nested).isInstanceOf(MapType.class); + } + + @Test + @DisplayName("toDataType exposes MAP with a NOT NULL key") + void toDataTypeExposesMap() { + assertThat(LanceTypeConverter.toDataType(mapOf(notNullString(), new IntType(true))).toString()) + .startsWith("MAP<") + .contains("NOT NULL"); + } + + @Test + @DisplayName("A value type the map path cannot move is rejected at DDL time, not on first write") + void unsupportedValueTypeIsRejectedAtDdlTime() { + // Arrow itself is happy with Map, so without this check CREATE TABLE would + // succeed and the failure would only surface inside the converter on the first write -- + // the same trap the Lance 2.2 format requirement set for MAP in the first place. + assertThatThrownBy( + () -> + LanceTypeConverter.flinkTypeToArrowField( + "attrs", + mapOf( + notNullString(), + new org.apache.flink.table.types.logical.DateType( + true)))) + .isInstanceOf(LanceTypeConverter.UnsupportedTypeException.class) + .hasMessageContaining("attrs") + .hasMessageContaining("MAP value"); + } + + @Test + @DisplayName("An unsupported key type is rejected the same way") + void unsupportedKeyTypeIsRejected() { + assertThatThrownBy( + () -> + LanceTypeConverter.flinkTypeToArrowField( + "attrs", + mapOf( + new org.apache.flink.table.types.logical.DateType( + false), + new IntType(true)))) + .isInstanceOf(LanceTypeConverter.UnsupportedTypeException.class) + .hasMessageContaining("MAP key"); + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterMapTest.java b/src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterMapTest.java new file mode 100644 index 0000000..edc898e --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/converter/RowDataConverterMapTest.java @@ -0,0 +1,171 @@ +/* + * 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.converter; + +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.MapData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Round-trips MAP columns through {@link RowDataConverter}. + * + *

    The schema mapping alone is not enough to make MAP usable: the read and write dispatches both + * had only a {@code ListVector} branch, and because {@code MapVector extends ListVector} a map + * column would have fallen into it and been handled with list semantics. + */ +class RowDataConverterMapTest { + + private static final RowType ROW_TYPE = + RowType.of( + new LogicalType[] { + new IntType(false), + new MapType( + true, + new VarCharType(false, VarCharType.MAX_LENGTH), + new IntType(true)) + }, + new String[] {"id", "attrs"}); + + private static MapData map(Object... kv) { + Map m = new HashMap<>(); + for (int i = 0; i < kv.length; i += 2) { + m.put(StringData.fromString((String) kv[i]), kv[i + 1]); + } + return new GenericMapData(m); + } + + private static List roundTrip(List input) { + RowDataConverter converter = new RowDataConverter(ROW_TYPE); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = converter.createVectorSchemaRoot(allocator)) { + converter.toVectorSchemaRoot(input, root); + return new ArrayList<>(converter.toRowDataList(root)); + } + } + + @Test + @DisplayName("A populated map survives the round-trip with its entries intact") + void populatedMapRoundTrips() { + List out = + roundTrip( + Arrays.asList( + GenericRowData.of(1, map("a", 10, "b", 20)))); + + assertThat(out).hasSize(1); + MapData read = out.get(0).getMap(1); + assertThat(read.size()).isEqualTo(2); + + Map asJava = new HashMap<>(); + for (int i = 0; i < read.size(); i++) { + asJava.put( + read.keyArray().getString(i).toString(), + read.valueArray().isNullAt(i) ? null : read.valueArray().getInt(i)); + } + assertThat(asJava).containsEntry("a", 10).containsEntry("b", 20); + } + + @Test + @DisplayName("An empty map stays empty and does not become NULL") + void emptyMapStaysEmpty() { + // Distinguishing the two matters: an entries struct left without its validity bit reads + // back as NULL, which would look like data loss rather than a missing flag. + List out = roundTrip(Arrays.asList(GenericRowData.of(1, map()))); + + assertThat(out.get(0).isNullAt(1)).isFalse(); + assertThat(out.get(0).getMap(1).size()).isZero(); + } + + @Test + @DisplayName("A NULL map stays NULL") + void nullMapStaysNull() { + List out = roundTrip(Arrays.asList(GenericRowData.of(1, null))); + + assertThat(out.get(0).isNullAt(1)).isTrue(); + } + + @Test + @DisplayName("Rows keep their own maps when populated, empty and NULL are interleaved") + void interleavedRowsKeepTheirOwnMaps() { + // Offsets are per row, so a mistake there leaks one row's entries into the next -- which a + // single-row test cannot see. + List out = + roundTrip( + Arrays.asList( + GenericRowData.of(1, map("x", 1)), + GenericRowData.of(2, map()), + GenericRowData.of(3, null), + GenericRowData.of(4, map("y", 2, "z", 3)))); + + assertThat(out).hasSize(4); + assertThat(out.get(0).getMap(1).size()).isEqualTo(1); + assertThat(out.get(0).getMap(1).keyArray().getString(0).toString()).isEqualTo("x"); + + assertThat(out.get(1).isNullAt(1)).isFalse(); + assertThat(out.get(1).getMap(1).size()).isZero(); + + assertThat(out.get(2).isNullAt(1)).isTrue(); + + assertThat(out.get(3).getMap(1).size()).isEqualTo(2); + } + + @Test + @DisplayName("A NULL map value is preserved, only the key may not be NULL") + void nullValueIsPreserved() { + Map m = new HashMap<>(); + m.put(StringData.fromString("k"), null); + List out = + roundTrip(Arrays.asList(GenericRowData.of(1, new GenericMapData(m)))); + + MapData read = out.get(0).getMap(1); + assertThat(read.size()).isEqualTo(1); + assertThat(read.valueArray().isNullAt(0)).isTrue(); + } + + @Test + @DisplayName("A NULL key is rejected rather than written as a broken entry") + void nullKeyIsRejected() { + Map m = new HashMap<>(); + m.put(null, 1); + + assertThatThrownBy( + () -> roundTrip(Arrays.asList(GenericRowData.of(1, new GenericMapData(m))))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("MAP key must not be NULL"); + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java b/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java index 6ac5bfc..3edee35 100644 --- a/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java +++ b/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java @@ -44,6 +44,7 @@ import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; import static org.assertj.core.api.Assertions.assertThatThrownBy; /** @@ -166,18 +167,49 @@ void testUnresolvedDataTypeRejected() { @Test @DisplayName("An unsupported column type is rejected as a schema problem") void testUnsupportedTypeRejected() { - // MAP has no Arrow mapping yet. DECIMAL used to stand in here, so this case had to move - // when DECIMAL gained one; the assertion is about how an unmappable type is surfaced, not - // about MAP specifically. + // MULTISET has no Arrow mapping yet. DECIMAL stood in here first, then MAP; each moved on + // as it gained a mapping. The assertion is about how an unmappable type is surfaced, not + // about MULTISET specifically. CatalogTable table = tableWith( Schema.newBuilder() - .column("attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.INT())) + .column("tags", DataTypes.MULTISET(DataTypes.STRING())) .build()); assertThatThrownBy(() -> LanceNamespaceCatalog.toArrowIpcSchema(table, allocator)) .isInstanceOf(org.apache.flink.table.catalog.exceptions.CatalogException.class) .hasMessageContaining("Cannot create a Lance table with this schema") - .hasMessageContaining("MapType"); + .hasMessageContaining("MultisetType"); + } + + @Test + @DisplayName("A MAP declared the obvious way is rejected, because its key is nullable") + void testMapWithNullableKeyRejected() { + // DataTypes.MAP(STRING(), INT()) produces a nullable key, which is the natural thing to + // write and which Arrow does not allow. The message has to say so explicitly, otherwise + // the user has no way to know that .notNull() on the key is what is missing. + CatalogTable table = tableWith( + Schema.newBuilder() + .column("attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.INT())) + .build()); + + assertThatThrownBy(() -> LanceNamespaceCatalog.toArrowIpcSchema(table, allocator)) + .isInstanceOf(org.apache.flink.table.catalog.exceptions.CatalogException.class) + .hasMessageContaining("attrs") + .hasMessageContaining("NOT NULL"); + } + + @Test + @DisplayName("A MAP with a NOT NULL key is accepted") + void testMapWithNotNullKeyAccepted() { + CatalogTable table = tableWith( + Schema.newBuilder() + .column( + "attrs", + DataTypes.MAP(DataTypes.STRING().notNull(), DataTypes.INT())) + .build()); + + assertThatCode(() -> LanceNamespaceCatalog.toArrowIpcSchema(table, allocator)) + .doesNotThrowAnyException(); } @Test From c32c7cc7048ce060a7fc41d0cea93eda141fdf1a Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 16:16:06 +0800 Subject: [PATCH 10/11] feat: support MULTISET columns on the MAP encoding Closes the last remaining type gap. A MULTISET is physically MAP and Flink already hands it around as MapData at runtime -- getDefaultConversion is java.util.Map and RowData.getMap works on it -- so it reuses the map path instead of getting a parallel implementation. readMap and writeMap now take the key and value LogicalTypes directly rather than a MapType, which is the whole of the sharing; the count side is pinned to a non-null INT32 because an element that is present has an occurrence count by definition. The storage requirement was carried over as a guess and is now measured. A MULTISET-shaped map fails to write on format 2.1 with the same error MAP hit -- "Map data type is only supported in Lance file format 2.2+", lance-encoding/src/encoder.rs:485 -- and writes cleanly on 2.2. Because the version check in LanceCatalog keys off ArrowType.Map, and setNull already has a MapVector branch ahead of ListVector, neither needed changing. A MULTISET reads back as MAP. An Arrow map carries nothing that separates MAP from MULTISET, and MAP is the far more common declaration, so an untagged map resolves to MAP. Field metadata was tried and does survive a Lance round-trip, so this was a choice rather than a limitation: it would put a Flink-specific key into a schema other engines also read, for a type that is mostly an aggregation result rather than a column declaration. Writes and stored bytes are unaffected; only the recovered type name differs. The NULL-key rejection message is now parameterised, since a user who wrote a MULTISET has no "key" in their DDL to go looking for. Same reason the element type check reports "MULTISET element" rather than "MAP key". The unsupported-type sample in LanceNamespaceCatalogSchemaTest moves a fourth time, to INTERVAL. DECIMAL, MAP and MULTISET were all types Lance genuinely stores, so each was always going to gain a mapping; an interval is not analytical storage data and is not on that path. --- ...sue-map-type-blocked-by-storage-version.md | 32 ++- docs/src/config.md | 28 +++ .../lance/converter/LanceTypeConverter.java | 63 ++++- .../lance/converter/RowDataConverter.java | 70 ++++-- .../MultisetColumnLanceRoundTripITCase.java | 183 ++++++++++++++ .../LanceTypeConverterMultisetTest.java | 226 ++++++++++++++++++ .../LanceNamespaceCatalogSchemaTest.java | 39 ++- 7 files changed, 617 insertions(+), 24 deletions(-) create mode 100644 src/test/java/org/apache/flink/connector/lance/MultisetColumnLanceRoundTripITCase.java create mode 100644 src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMultisetTest.java diff --git a/.gh-comments/issue-map-type-blocked-by-storage-version.md b/.gh-comments/issue-map-type-blocked-by-storage-version.md index aef049f..f649ac8 100644 --- a/.gh-comments/issue-map-type-blocked-by-storage-version.md +++ b/.gh-comments/issue-map-type-blocked-by-storage-version.md @@ -108,4 +108,34 @@ DECIMAL、后用 MAP，两者各自获得映射后都得换；现已改用 `MULT - 2.1 与 2.2 数据集能否在同一路径混用（时间旅行读旧版本）。 - 2.2 是否影响其他类型的编码或既有索引。 -- `MULTISET`（Flink 侧是 `Map`）仍无映射，推测同受 2.2 限制，未实测。 + +## MULTISET —— 已完成 + +推测得到证实：MULTISET 形状的 map（`Map`）在 2.1 上建表成功、写入失败，报的 +是与 MAP 完全相同的 `Map data type is only supported in Lance file format 2.2+ +(lance-encoding/src/encoder.rs:485)`。2.2 上写入正常。 + +实现完全复用 MAP 的通路：探针确认 `MultisetType.getDefaultConversion()` 是 +`java.util.Map`、`RowData.getMap()` 可用，因此 MULTISET 在运行时就是 `MapData`。 +`readMap`/`writeMap` 的签名从收 `MapType` 改为直接收 key/value 两个 `LogicalType`，MAP 与 +MULTISET 共用同一实现，不存在平行分支。count 侧固定为非空 INT32。 + +`LanceCatalog` 的版本校验无需改动——`containsMap` 按 `ArrowType.Map` 判断，MULTISET 映射成 +同一 Arrow 类型，自动覆盖。`setNull` 同理，上一轮加的 MapVector 前置分支直接生效。 + +### 一个有意接受的不对称 + +MULTISET 列**读回后呈现为 `MAP`**。Arrow map 不携带任何可区分 +`MAP` 与 `MULTISET` 的信息，而 MAP 是远更常见的声明，因此无 +标记的 map 一律解析为 MAP。 + +我实测过打元数据的可行性：`FieldType` 支持自定义 metadata，且 Lance **确实**原样往返了 +`{flink.type=MULTISET}`。技术上可行，但我选择不用——这会为一个极少作为表列声明的类型 +（MULTISET 主要来自 `COLLECT()` 聚合结果）在跨引擎共读的 schema 里塞进一个 `flink.` 专有 +键，正是 A7 遗留项在推动消除的那类耦合。写入与存储数据不受影响，仅恢复出的类型名不同。 + +### 「不支持类型」测试样本第四次搬家，这次应该是最后一次 + +`LanceNamespaceCatalogSchemaTest` 的样本历经 DECIMAL → MAP → MULTISET，现改为 +`INTERVAL`。前三者都是 Lance 真正会存储的数据类型，所以迟早都会获得映射；区间类型不属于 +分析存储的数据，不在这条路径上。 diff --git a/docs/src/config.md b/docs/src/config.md index 7f2ae18..2d2c103 100644 --- a/docs/src/config.md +++ b/docs/src/config.md @@ -67,6 +67,34 @@ narrower than what a top-level column accepts — the same limit applies to than on the first write. An empty map and a `NULL` map are stored distinctly. A `NULL` value is allowed; a `NULL` key is not. +### MULTISET columns + +A `MULTISET` is stored as `MAP`, so it carries the same 2.2 +requirement, the same element type restrictions, and the same `NOT NULL` rule — +the element becomes the map key: + +```sql +-- Rejected: DataTypes.MULTISET(DataTypes.STRING()) yields a nullable element. +CREATE TABLE t (tags MULTISET) WITH (...); + +-- Correct. +CREATE TABLE t (tags MULTISET) WITH ( + 'write.data-storage-version' = '2.2', ... +); +``` + +The count side is always a non-null `INT`, since an element that is present has +an occurrence count by definition. + +One asymmetry is worth knowing about: a `MULTISET` column **reads back as +`MAP`**. An Arrow map carries nothing that separates +`MAP` from `MULTISET`, and `MAP` is the far more +common declaration, so an untagged map resolves to `MAP`. Writes are unaffected, +and the stored data is identical either way — only the recovered type name +differs. Tagging the field with metadata would make the distinction survive, but +it would put a Flink-specific key into a schema other engines also read, so it +is deliberately not done. + ### Vector index | Option | Required | Default | Description | diff --git a/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java b/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java index 089bd13..41ee192 100644 --- a/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java +++ b/src/main/java/org/apache/flink/connector/lance/converter/LanceTypeConverter.java @@ -32,6 +32,7 @@ import org.apache.flink.table.types.logical.LocalZonedTimestampType; import org.apache.flink.table.types.logical.LogicalType; import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.MultisetType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.table.types.logical.SmallIntType; import org.apache.flink.table.types.logical.TimeType; @@ -320,6 +321,21 @@ public static Field flinkTypeToArrowField(String name, LogicalType logicalType) // keysSorted=false: nothing in the write path sorts entries, and claiming otherwise // would let a reader skip its own ordering work on unordered data. arrowType = new ArrowType.Map(false); + } else if (logicalType instanceof MultisetType) { + // A MULTISET is physically a MAP, which is also how Flink represents it + // at runtime (getDefaultConversion is java.util.Map, and RowData.getMap works on it). + // So it reuses the map encoding wholesale rather than getting its own. + MultisetType multisetType = (MultisetType) logicalType; + LogicalType elementType = multisetType.getElementType(); + if (elementType.isNullable()) { + throw new UnsupportedTypeException( + "MULTISET element must be NOT NULL for column '" + name + "': the element " + + "becomes an Arrow map key, which cannot be nullable. Declare it " + + "as e.g. MULTISET."); + } + children = new ArrayList<>(); + children.add(multisetEntriesField(name, elementType)); + arrowType = new ArrowType.Map(false); } else if (logicalType instanceof RowType) { RowType rowType = (RowType) logicalType; children = new ArrayList<>(); @@ -351,8 +367,8 @@ private static Field mapEntriesField( // Arrow accepts far more element types here than the read/write path can actually move, so // an unchecked MAP would create a table whose first write fails deep in the // converter. Reject it at DDL time instead, where the message can name the column. - requireSupportedMapElement(mapColumnName, "key", keyType); - requireSupportedMapElement(mapColumnName, "value", valueType); + requireSupportedMapElement(mapColumnName, "MAP key", keyType); + requireSupportedMapElement(mapColumnName, "MAP value", valueType); Field keyField = flinkTypeToArrowField(MapVector.KEY_NAME, keyType); if (keyField.isNullable()) { @@ -376,12 +392,49 @@ private static Field mapEntriesField( entryChildren); } + /** + * Build the {@code entries} struct backing a multiset. + * + *

    Identical in shape to {@link #mapEntriesField}, with the value side pinned to a non-null + * INT: a multiset's value is an occurrence count, never user data, so it is neither nullable nor + * variable in type. + */ + private static Field multisetEntriesField(String columnName, LogicalType elementType) { + requireSupportedMapElement(columnName, "MULTISET element", elementType); + + Field keyField = flinkTypeToArrowField(MapVector.KEY_NAME, elementType); + if (keyField.isNullable()) { + keyField = + new Field( + MapVector.KEY_NAME, + new FieldType(false, keyField.getType(), null), + keyField.getChildren()); + } + Field countField = + new Field( + MapVector.VALUE_NAME, + new FieldType(false, new ArrowType.Int(32, true), null), + null); + + List entryChildren = new ArrayList<>(); + entryChildren.add(keyField); + entryChildren.add(countField); + + return new Field( + MapVector.DATA_VECTOR_NAME, + new FieldType(false, ArrowType.Struct.INSTANCE, null), + entryChildren); + } + /** * Element types the map read/write path can carry. * *

    This mirrors what {@code RowDataConverter}'s array element helpers implement, since map * keys and values reuse them. It is narrower than the set of types allowed for a top-level * column, and widening it means extending those helpers first. + * + * @param role full description such as {@code "MAP key"} or {@code "MULTISET element"}; it is + * used verbatim so the message matches the DDL the user actually wrote */ private static void requireSupportedMapElement( String mapColumnName, String role, LogicalType elementType) { @@ -393,7 +446,7 @@ private static void requireSupportedMapElement( || elementType instanceof VarCharType; if (!supported) { throw new UnsupportedTypeException( - "Unsupported MAP " + "Unsupported " + role + " type for column '" + mapColumnName @@ -544,6 +597,10 @@ public static DataType toDataType(LogicalType logicalType) { DataType valueDataType = toDataType(mapType.getValueType()); // The key stays NOT NULL to match Arrow, which does not allow a nullable map key. return DataTypes.MAP(keyDataType.notNull(), valueDataType); + } else if (logicalType instanceof MultisetType) { + MultisetType multisetType = (MultisetType) logicalType; + // The element becomes the map key, so it carries the same NOT NULL requirement. + return DataTypes.MULTISET(toDataType(multisetType.getElementType()).notNull()); } else if (logicalType instanceof RowType) { RowType rowType = (RowType) logicalType; DataTypes.Field[] fields = rowType.getFields().stream() diff --git a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java index 5059c8c..6794825 100644 --- a/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java +++ b/src/main/java/org/apache/flink/connector/lance/converter/RowDataConverter.java @@ -39,6 +39,7 @@ import org.apache.flink.table.types.logical.LocalZonedTimestampType; import org.apache.flink.table.types.logical.LogicalType; import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.MultisetType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.table.types.logical.SmallIntType; import org.apache.flink.table.types.logical.TimeType; @@ -99,6 +100,12 @@ */ public class RowDataConverter implements Serializable { + /** + * The value side of a multiset is an occurrence count: always present, always an int. Pinning it + * here keeps the read and write paths from disagreeing about it. + */ + private static final LogicalType MULTISET_COUNT_TYPE = new IntType(false); + private static final long serialVersionUID = 1L; private static final Logger LOG = LoggerFactory.getLogger(RowDataConverter.class); @@ -231,7 +238,13 @@ private Object readValue(FieldVector vector, int index, LogicalType logicalType) } else if (logicalType instanceof ArrayType) { return readArray(vector, index, (ArrayType) logicalType); } else if (logicalType instanceof MapType) { - return readMap(vector, index, (MapType) logicalType); + MapType mapType = (MapType) logicalType; + return readMap(vector, index, mapType.getKeyType(), mapType.getValueType()); + } else if (logicalType instanceof MultisetType) { + // A multiset is stored as MAP, so it reads back through the same path + // with the count side pinned to a non-null INT. + MultisetType multisetType = (MultisetType) logicalType; + return readMap(vector, index, multisetType.getElementType(), MULTISET_COUNT_TYPE); } else if (logicalType instanceof RowType) { return readStruct(vector, index, (RowType) logicalType); } @@ -346,7 +359,8 @@ private ArrayData readArray(FieldVector vector, int index, ArrayType arrayType) * and a {@code value} child. The offsets come from the enclosing list, so the key and value * slices are read from the same index range. */ - private MapData readMap(FieldVector vector, int index, MapType mapType) { + private MapData readMap( + FieldVector vector, int index, LogicalType keyType, LogicalType valueType) { if (!(vector instanceof MapVector)) { // Note this cannot be relaxed to ListVector: a plain list carries no key child, and // treating one as a map would read garbage out of the element vector. @@ -368,9 +382,9 @@ private MapData readMap(FieldVector vector, int index, MapType mapType) { + MapVector.VALUE_NAME + "' children"); } - ArrayData keys = readArrayData(keyVector, startIndex, size, mapType.getKeyType()); - ArrayData values = readArrayData(valueVector, startIndex, size, mapType.getValueType()); - return new GenericMapData(toJavaMap(keys, values, mapType)); + ArrayData keys = readArrayData(keyVector, startIndex, size, keyType); + ArrayData values = readArrayData(valueVector, startIndex, size, valueType); + return new GenericMapData(toJavaMap(keys, values, keyType, valueType)); } /** @@ -379,11 +393,12 @@ private MapData readMap(FieldVector vector, int index, MapType mapType) { *

    Duplicate keys collapse to the last occurrence, matching how Flink's own map * implementations behave when a duplicate reaches them. */ - private Map toJavaMap(ArrayData keys, ArrayData values, MapType mapType) { + private Map toJavaMap( + ArrayData keys, ArrayData values, LogicalType keyType, LogicalType valueType) { Map result = new LinkedHashMap<>(); for (int i = 0; i < keys.size(); i++) { - Object key = elementAt(keys, i, mapType.getKeyType()); - Object value = elementAt(values, i, mapType.getValueType()); + Object key = elementAt(keys, i, keyType); + Object value = elementAt(values, i, valueType); result.put(key, value); } return result; @@ -525,7 +540,8 @@ private Object getFieldValue(RowData rowData, int index, LogicalType logicalType return rowData.getDecimal(index, decimalType.getPrecision(), decimalType.getScale()); } else if (logicalType instanceof ArrayType) { return rowData.getArray(index); - } else if (logicalType instanceof MapType) { + } else if (logicalType instanceof MapType || logicalType instanceof MultisetType) { + // Flink hands both back as MapData; MultisetType.getDefaultConversion is java.util.Map. return rowData.getMap(index); } else if (logicalType instanceof RowType) { RowType nestedRowType = (RowType) logicalType; @@ -579,7 +595,23 @@ private void writeValue(FieldVector vector, int index, Object value, LogicalType } else if (logicalType instanceof ArrayType) { writeArray(vector, index, (ArrayData) value, (ArrayType) logicalType); } else if (logicalType instanceof MapType) { - writeMap(vector, index, (MapData) value, (MapType) logicalType); + MapType mapType = (MapType) logicalType; + writeMap( + vector, + index, + (MapData) value, + mapType.getKeyType(), + mapType.getValueType(), + "MAP key"); + } else if (logicalType instanceof MultisetType) { + MultisetType multisetType = (MultisetType) logicalType; + writeMap( + vector, + index, + (MapData) value, + multisetType.getElementType(), + MULTISET_COUNT_TYPE, + "MULTISET element"); } else if (logicalType instanceof RowType) { writeStruct(vector, index, (RowData) value, (RowType) logicalType); } else { @@ -776,7 +808,13 @@ private void writeArray(FieldVector vector, int index, ArrayData arrayData, Arra * even though key and value were written, and a NULL key is rejected because Arrow does not * allow one. */ - private void writeMap(FieldVector vector, int index, MapData mapData, MapType mapType) { + private void writeMap( + FieldVector vector, + int index, + MapData mapData, + LogicalType keyType, + LogicalType valueType, + String keyRole) { if (!(vector instanceof MapVector)) { throw new LanceTypeConverter.UnsupportedTypeException( "Unsupported map Vector type: " + vector.getClass().getSimpleName()); @@ -789,9 +827,11 @@ private void writeMap(FieldVector vector, int index, MapData mapData, MapType ma for (int i = 0; i < size; i++) { if (keys.isNullAt(i)) { + // keyRole names what the user actually wrote -- a MULTISET has no "key", so + // reporting one would send them looking for something that is not in their DDL. throw new IllegalArgumentException( - "MAP key must not be NULL: Arrow map keys are non-nullable, so a NULL key " - + "cannot be written (entry " + i + ")"); + keyRole + " must not be NULL: Arrow map keys are non-nullable, so a NULL " + + "key cannot be written (entry " + i + ")"); } } @@ -802,8 +842,8 @@ private void writeMap(FieldVector vector, int index, MapData mapData, MapType ma FieldVector valueVector = entries.getChild(MapVector.VALUE_NAME); int startIndex = mapVector.getElementStartIndex(index); - writeArrayData(keyVector, startIndex, keys, mapType.getKeyType()); - writeArrayData(valueVector, startIndex, values, mapType.getValueType()); + writeArrayData(keyVector, startIndex, keys, keyType); + writeArrayData(valueVector, startIndex, values, valueType); // Without this the entries struct keeps a zero validity bit and the whole entry reads back // as NULL, which looks like data loss rather than a missing flag. diff --git a/src/test/java/org/apache/flink/connector/lance/MultisetColumnLanceRoundTripITCase.java b/src/test/java/org/apache/flink/connector/lance/MultisetColumnLanceRoundTripITCase.java new file mode 100644 index 0000000..ef53c33 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/MultisetColumnLanceRoundTripITCase.java @@ -0,0 +1,183 @@ +/* + * 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.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.flink.connector.lance.converter.LanceTypeConverter; +import org.apache.flink.connector.lance.converter.RowDataConverter; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.MapData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MultisetType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.lance.Dataset; +import org.lance.ReadOptions; +import org.lance.WriteParams; +import org.lance.merge.MergeInsertParams; + +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * End-to-end MULTISET support against a real Lance dataset. + * + *

    The in-memory converter test reads back the vectors it just wrote and never passes through + * Lance's own validation, so it cannot see a non-null constraint violation. Only a real dataset can + * — which is how the entries-struct validity bug was caught for MAP. + */ +class MultisetColumnLanceRoundTripITCase { + + @TempDir Path tempDir; + + private static final RowType ROW_TYPE = + RowType.of( + new LogicalType[] { + new IntType(false), + new MultisetType(true, new VarCharType(false, VarCharType.MAX_LENGTH)) + }, + new String[] {"id", "tags"}); + + private static MapData counts(Object... kv) { + Map m = new HashMap<>(); + for (int i = 0; i < kv.length; i += 2) { + m.put(StringData.fromString((String) kv[i]), kv[i + 1]); + } + return new GenericMapData(m); + } + + @Test + @DisplayName("MULTISET rows survive a real Lance round-trip with their counts") + void multisetSurvivesLanceRoundTrip() throws Exception { + String path = tempDir.resolve("multiset_rows").toString(); + + List input = + Arrays.asList( + GenericRowData.of(1, counts("a", 2, "b", 1)), + GenericRowData.of(2, counts()), + GenericRowData.of(3, null), + GenericRowData.of(4, counts("solo", 5))); + + RowDataConverter converter = new RowDataConverter(ROW_TYPE); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + // A multiset is stored as a map, so it inherits the Lance 2.2 format requirement. + Dataset.create( + allocator, + path, + LanceTypeConverter.toArrowSchema(ROW_TYPE), + new WriteParams.Builder().withDataStorageVersion("2.2").build()) + .close(); + + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build()); + VectorSchemaRoot root = converter.createVectorSchemaRoot(allocator)) { + converter.toVectorSchemaRoot(input, root); + MergeInsertParams params = + new MergeInsertParams(Collections.singletonList("id")) + .withMatchedUpdateAll() + .withNotMatched(MergeInsertParams.WhenNotMatched.InsertAll); + ArrowArrayStreamsTestAccess.mergeInsert(ds, params, allocator, root); + } + + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build())) { + assertThat(ds.countRows()).isEqualTo(4L); + + List readBack = new ArrayList<>(); + try (org.apache.arrow.vector.ipc.ArrowReader reader = ds.newScan().scanBatches()) { + while (reader.loadNextBatch()) { + readBack.addAll(converter.toRowDataList(reader.getVectorSchemaRoot())); + } + } + + // mergeInsert does not preserve input order, so key the rows by id. + Map byId = new HashMap<>(); + for (RowData row : readBack) { + byId.put(row.getInt(0), row); + } + assertThat(byId).hasSize(4); + + assertThat(flatten(byId.get(1).getMap(1))) + .containsEntry("a", 2) + .containsEntry("b", 1); + + assertThat(byId.get(2).isNullAt(1)).isFalse(); + assertThat(byId.get(2).getMap(1).size()).isZero(); + + assertThat(byId.get(3).isNullAt(1)).isTrue(); + + assertThat(flatten(byId.get(4).getMap(1))) + .containsExactlyEntriesOf(Collections.singletonMap("solo", 5)); + } + } + } + + @Test + @DisplayName("The persisted schema stores a MULTISET as a map with a non-null int count") + void persistedSchemaShapeIsAMap() throws Exception { + String path = tempDir.resolve("multiset_schema").toString(); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Dataset.create( + allocator, + path, + LanceTypeConverter.toArrowSchema(ROW_TYPE), + new WriteParams.Builder().withDataStorageVersion("2.2").build()) + .close(); + + try (Dataset ds = Dataset.open(allocator, path, new ReadOptions.Builder().build())) { + org.apache.arrow.vector.types.pojo.Field tags = ds.getSchema().findField("tags"); + assertThat(tags.getType()) + .isInstanceOf(org.apache.arrow.vector.types.pojo.ArrowType.Map.class); + + org.apache.arrow.vector.types.pojo.Field entries = tags.getChildren().get(0); + assertThat(entries.isNullable()).isFalse(); + + org.apache.arrow.vector.types.pojo.Field count = entries.getChildren().get(1); + assertThat(count.getType()) + .isInstanceOf(org.apache.arrow.vector.types.pojo.ArrowType.Int.class); + // Lance has to agree the count is non-null; if it were declared nullable the + // writer would accept the schema and then reject the first row. + assertThat(count.isNullable()).isFalse(); + } + } + } + + private static Map flatten(MapData mapData) { + Map out = new HashMap<>(); + for (int i = 0; i < mapData.size(); i++) { + out.put(mapData.keyArray().getString(i).toString(), mapData.valueArray().getInt(i)); + } + return out; + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMultisetTest.java b/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMultisetTest.java new file mode 100644 index 0000000..1633c56 --- /dev/null +++ b/src/test/java/org/apache/flink/connector/lance/converter/LanceTypeConverterMultisetTest.java @@ -0,0 +1,226 @@ +/* + * 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.converter; + +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.flink.table.data.GenericMapData; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.MapData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; +import org.apache.flink.table.types.logical.MultisetType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * MULTISET support, which is layered on the MAP encoding. + * + *

    A multiset is physically {@code MAP} and Flink hands it around as + * {@link MapData} at runtime, so it reuses the map read/write path rather than getting its own. + */ +class LanceTypeConverterMultisetTest { + + private static LogicalType notNullString() { + return new VarCharType(false, VarCharType.MAX_LENGTH); + } + + private static MultisetType multisetOf(LogicalType element) { + return new MultisetType(true, element); + } + + @Test + @DisplayName("A MULTISET becomes an Arrow map whose value side is a non-null INT count") + void multisetBecomesArrowMapWithIntCount() { + Field field = + LanceTypeConverter.flinkTypeToArrowField("tags", multisetOf(notNullString())); + + assertThat(field.getType()).isInstanceOf(ArrowType.Map.class); + Field entries = field.getChildren().get(0); + assertThat(entries.getName()).isEqualTo("entries"); + assertThat(entries.isNullable()).isFalse(); + + Field key = entries.getChildren().get(0); + Field count = entries.getChildren().get(1); + assertThat(key.getName()).isEqualTo(MapVector.KEY_NAME); + assertThat(key.isNullable()).isFalse(); + + // The count is an occurrence tally, never user data: always present, always an int. + assertThat(count.getName()).isEqualTo(MapVector.VALUE_NAME); + assertThat(count.getType()).isInstanceOf(ArrowType.Int.class); + assertThat(((ArrowType.Int) count.getType()).getBitWidth()).isEqualTo(32); + assertThat(count.isNullable()) + .as("a count cannot be absent for an element that is present") + .isFalse(); + } + + @Test + @DisplayName("A nullable MULTISET element is rejected, naming the column") + void nullableElementIsRejected() { + // The element becomes the Arrow map key, and Arrow keys cannot be nullable. This matters + // because DataTypes.MULTISET(DataTypes.STRING()) produces exactly this. + assertThatThrownBy( + () -> + LanceTypeConverter.flinkTypeToArrowField( + "tags", + multisetOf(new VarCharType(true, VarCharType.MAX_LENGTH)))) + .isInstanceOf(LanceTypeConverter.UnsupportedTypeException.class) + .hasMessageContaining("tags") + .hasMessageContaining("MULTISET element") + .hasMessageContaining("NOT NULL"); + } + + @Test + @DisplayName("An element type the map path cannot move is rejected at DDL time") + void unsupportedElementTypeIsRejected() { + assertThatThrownBy( + () -> + LanceTypeConverter.flinkTypeToArrowField( + "tags", + multisetOf( + new org.apache.flink.table.types.logical.DateType( + false)))) + .isInstanceOf(LanceTypeConverter.UnsupportedTypeException.class) + // The message must say MULTISET element, not "MAP key" -- a user who wrote a + // MULTISET has no key in their DDL to go looking for. + .hasMessageContaining("MULTISET element") + .hasMessageContaining("tags"); + } + + @Test + @DisplayName("A MULTISET reads back as MAP, because the two are physically identical") + void multisetReadsBackAsMap() { + // Deliberate trade-off. An Arrow map carries nothing that distinguishes MAP + // from MULTISET, and MAP is by far the more common declaration, so an untagged + // map resolves to MAP. Tagging the field with metadata would work -- Lance does round-trip + // field metadata -- but it would add a Flink-specific key to the stored schema for the sake + // of a rarely declared type, which is the kind of cross-engine coupling being removed + // elsewhere. + Field field = + LanceTypeConverter.flinkTypeToArrowField("tags", multisetOf(notNullString())); + + LogicalType back = LanceTypeConverter.arrowTypeToFlinkType(field); + + assertThat(back).isInstanceOf(MapType.class); + assertThat(back).isNotInstanceOf(MultisetType.class); + MapType asMap = (MapType) back; + assertThat(asMap.getKeyType()).isInstanceOf(VarCharType.class); + assertThat(asMap.getValueType()).isInstanceOf(IntType.class); + } + + @Test + @DisplayName("toDataType exposes MULTISET with a NOT NULL element") + void toDataTypeExposesMultiset() { + assertThat(LanceTypeConverter.toDataType(multisetOf(notNullString())).toString()) + .startsWith("MULTISET<") + .contains("NOT NULL"); + } + + @Test + @DisplayName("Counts survive a round-trip through the converter") + void countsRoundTrip() { + RowType rowType = + RowType.of( + new LogicalType[] {new IntType(false), multisetOf(notNullString())}, + new String[] {"id", "tags"}); + + Map counts = new HashMap<>(); + counts.put(StringData.fromString("a"), 2); + counts.put(StringData.fromString("b"), 1); + + List out = + roundTrip(rowType, Arrays.asList(GenericRowData.of(1, new GenericMapData(counts)))); + + MapData read = out.get(0).getMap(1); + assertThat(read.size()).isEqualTo(2); + + Map asJava = new HashMap<>(); + for (int i = 0; i < read.size(); i++) { + asJava.put(read.keyArray().getString(i).toString(), read.valueArray().getInt(i)); + } + assertThat(asJava).containsEntry("a", 2).containsEntry("b", 1); + } + + @Test + @DisplayName("An empty MULTISET stays distinct from a NULL one") + void emptyAndNullAreDistinct() { + RowType rowType = + RowType.of( + new LogicalType[] {new IntType(false), multisetOf(notNullString())}, + new String[] {"id", "tags"}); + + List out = + roundTrip( + rowType, + Arrays.asList( + GenericRowData.of(1, new GenericMapData(new HashMap<>())), + GenericRowData.of(2, null))); + + assertThat(out.get(0).isNullAt(1)).isFalse(); + assertThat(out.get(0).getMap(1).size()).isZero(); + assertThat(out.get(1).isNullAt(1)).isTrue(); + } + + @Test + @DisplayName("A NULL element is rejected with MULTISET wording, not MAP wording") + void nullElementIsRejectedWithMultisetWording() { + RowType rowType = + RowType.of( + new LogicalType[] {new IntType(false), multisetOf(notNullString())}, + new String[] {"id", "tags"}); + + Map counts = new HashMap<>(); + counts.put(null, 1); + + assertThatThrownBy( + () -> + roundTrip( + rowType, + Arrays.asList( + GenericRowData.of(1, new GenericMapData(counts))))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("MULTISET element"); + } + + private static List roundTrip(RowType rowType, List input) { + RowDataConverter converter = new RowDataConverter(rowType); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot root = converter.createVectorSchemaRoot(allocator)) { + converter.toVectorSchemaRoot(input, root); + return new ArrayList<>(converter.toRowDataList(root)); + } + } +} diff --git a/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java b/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java index 3edee35..6ea971c 100644 --- a/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java +++ b/src/test/java/org/apache/flink/connector/lance/table/LanceNamespaceCatalogSchemaTest.java @@ -167,18 +167,19 @@ void testUnresolvedDataTypeRejected() { @Test @DisplayName("An unsupported column type is rejected as a schema problem") void testUnsupportedTypeRejected() { - // MULTISET has no Arrow mapping yet. DECIMAL stood in here first, then MAP; each moved on - // as it gained a mapping. The assertion is about how an unmappable type is surfaced, not - // about MULTISET specifically. + // INTERVAL has no Arrow mapping. This case has already moved three times -- DECIMAL, then + // MAP, then MULTISET -- because each was a type Lance genuinely stores and so each + // eventually gained a mapping. An interval is not analytical storage data, so it is not on + // that path and this sample should stop moving. CatalogTable table = tableWith( Schema.newBuilder() - .column("tags", DataTypes.MULTISET(DataTypes.STRING())) + .column("gap", DataTypes.INTERVAL(DataTypes.DAY())) .build()); assertThatThrownBy(() -> LanceNamespaceCatalog.toArrowIpcSchema(table, allocator)) .isInstanceOf(org.apache.flink.table.catalog.exceptions.CatalogException.class) .hasMessageContaining("Cannot create a Lance table with this schema") - .hasMessageContaining("MultisetType"); + .hasMessageContaining("DayTimeIntervalType"); } @Test @@ -212,6 +213,34 @@ void testMapWithNotNullKeyAccepted() { .doesNotThrowAnyException(); } + @Test + @DisplayName("A MULTISET declared the obvious way is rejected, because its element is nullable") + void testMultisetWithNullableElementRejected() { + // Same trap as MAP: DataTypes.MULTISET(DataTypes.STRING()) yields a nullable element, and + // the element becomes the Arrow map key, which cannot be nullable. + CatalogTable table = tableWith( + Schema.newBuilder() + .column("tags", DataTypes.MULTISET(DataTypes.STRING())) + .build()); + + assertThatThrownBy(() -> LanceNamespaceCatalog.toArrowIpcSchema(table, allocator)) + .isInstanceOf(org.apache.flink.table.catalog.exceptions.CatalogException.class) + .hasMessageContaining("tags") + .hasMessageContaining("NOT NULL"); + } + + @Test + @DisplayName("A MULTISET with a NOT NULL element is accepted") + void testMultisetWithNotNullElementAccepted() { + CatalogTable table = tableWith( + Schema.newBuilder() + .column("tags", DataTypes.MULTISET(DataTypes.STRING().notNull())) + .build()); + + assertThatCode(() -> LanceNamespaceCatalog.toArrowIpcSchema(table, allocator)) + .doesNotThrowAnyException(); + } + @Test @DisplayName("A resolved table uses its resolved physical schema") void testResolvedTableUsesPhysicalSchema() { From fcf713a88a8cc6eb32ddd942ede0e605fc141456 Mon Sep 17 00:00:00 2001 From: rockyyin Date: Tue, 22 Sep 2026 17:21:42 +0800 Subject: [PATCH 11/11] build: open java.util for Flink's Kryo serializer on JDK 17+ The batch-mode tests added with SupportsRowLevelUpdate go through Flink's SortingDataInput, which serialises via Kryo, which reflects into Arrays$ArrayList to read its backing array. JDK 17 refuses that without --add-opens and the job dies with InaccessibleObjectException; JDK 11 only warns, so a local run on 11 shows nothing and the gap surfaces only on the 17 and 21 legs of the CI matrix. Verified on 11, 17 and 21 across all three modules. --- pom.xml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 24ba54f..4912c0e 100644 --- a/pom.xml +++ b/pom.xml @@ -22,10 +22,16 @@ JDK 17 encapsulates that, so MemoryUtil's static initializer throws a RuntimeException naming this exact flag. On JDK 11 it is only an illegal-reflective-access warning, and passing the flag there is harmless. + + java.util is opened for Flink's Kryo serializer, which reflects into + Arrays$ArrayList to read its backing array. Batch-mode tests hit this through + SortingDataInput; on JDK 11 the reflection only warns, so the omission is + invisible there and only surfaces on 17+ as InaccessibleObjectException. + surefire and failsafe both read ${argLine} by default, and jacoco:prepare-agent appends its -javaagent to this property rather than replacing it, so the two coexist. --> - --add-opens=java.base/java.nio=ALL-UNNAMED + --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.util=ALL-UNNAMED 0.1.0