diff --git a/.java-version b/.java-version index 4099407..7273c0f 100644 --- a/.java-version +++ b/.java-version @@ -1 +1 @@ -23 +25 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 138978a..8ac5111 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -59,5 +59,5 @@ This starts a Trino server on port 8080 with the Lance connector configured. ### Requirements -- Java 23 or later -- Trino 476 or compatible version +- Java 25 or later +- Trino 481 or compatible version diff --git a/docs/src/install.md b/docs/src/install.md index 7ffb15e..753ba25 100644 --- a/docs/src/install.md +++ b/docs/src/install.md @@ -2,12 +2,12 @@ ## Requirements -- Java 23 or later -- Trino 476 or compatible version +- Java 25 or later +- Trino 481 or compatible version ## Download from GitHub Releases -Each release includes a `trino-lance--trino.tar.gz` archive containing all required JARs. Download from the [releases page](https://github.com/lancedb/lance-trino/releases). +Each release includes a `lance-trino--trino.tar.gz` archive containing all required JARs. Download from the [releases page](https://github.com/lancedb/lance-trino/releases). ### Quick Installation @@ -15,24 +15,24 @@ Each release includes a `trino-lance--trino.tar.gz` arch ```bash # Set variables VERSION="0.3.2" - TRINO_VERSION="476" + TRINO_VERSION="481" PLUGIN_DIR="/usr/lib/trino/plugin" # Download and extract - wget "https://github.com/lancedb/lance-trino/releases/download/v${VERSION}/trino-lance-${VERSION}-trino${TRINO_VERSION}.tar.gz" - tar -xzf "trino-lance-${VERSION}-trino${TRINO_VERSION}.tar.gz" -C "${PLUGIN_DIR}/" - mv "${PLUGIN_DIR}/trino-lance-${TRINO_VERSION}" "${PLUGIN_DIR}/lance" + wget "https://github.com/lancedb/lance-trino/releases/download/v${VERSION}/lance-trino-${VERSION}-trino${TRINO_VERSION}.tar.gz" + tar -xzf "lance-trino-${VERSION}-trino${TRINO_VERSION}.tar.gz" -C "${PLUGIN_DIR}/" + mv "${PLUGIN_DIR}/lance-trino-${VERSION}" "${PLUGIN_DIR}/lance" ``` === "Docker" ```dockerfile - FROM trinodb/trino:476 + FROM trinodb/trino:481 # Download and install Lance connector ARG VERSION=0.3.2 - ARG TRINO_VERSION=476 + ARG TRINO_VERSION=481 - RUN curl -fsSL "https://github.com/lancedb/lance-trino/releases/download/v${VERSION}/trino-lance-${VERSION}-trino${TRINO_VERSION}.tar.gz" \ + RUN curl -fsSL "https://github.com/lancedb/lance-trino/releases/download/v${VERSION}/lance-trino-${VERSION}-trino${TRINO_VERSION}.tar.gz" \ | tar -xz -C /usr/lib/trino/plugin/ \ - && mv "/usr/lib/trino/plugin/trino-lance-${TRINO_VERSION}" /usr/lib/trino/plugin/lance + && mv "/usr/lib/trino/plugin/lance-trino-${VERSION}" /usr/lib/trino/plugin/lance ``` diff --git a/plugin/trino-lance/pom.xml b/plugin/trino-lance/pom.xml index 632f9ed..1a37724 100755 --- a/plugin/trino-lance/pom.xml +++ b/plugin/trino-lance/pom.xml @@ -29,6 +29,7 @@ com.google.inject guice + classes @@ -108,6 +109,11 @@ arrow-memory-core + + org.apache.arrow + arrow-memory-unsafe + + org.apache.arrow arrow-vector @@ -172,12 +178,6 @@ provided - - org.openjdk.jol - jol-core - provided - - com.fasterxml.jackson.dataformat jackson-dataformat-yaml @@ -259,31 +259,12 @@ test - - io.trino - trino-spi - test-jar - test - - io.trino trino-testing test - - io.trino - trino-testing-kafka - test - - - org.glassfish.jersey.core - jersey-common - - - - io.trino trino-tpch @@ -328,19 +309,19 @@ org.testcontainers - kafka + testcontainers test org.testcontainers - localstack + testcontainers-kafka test org.testcontainers - testcontainers + testcontainers-localstack test @@ -401,7 +382,7 @@ - + io.trino:trino-spi io.trino:trino-matching io.trino:trino-cache @@ -414,7 +395,7 @@ maven-surefire-plugin - --add-opens=java.base/java.nio=ALL-UNNAMED + ${air.test.jvm.additional-arguments} --add-opens=java.base/java.nio=ALL-UNNAMED false ${surefire.forkCount} + Unsafe same_thread @@ -435,6 +417,8 @@ com.google.protobuf:protobuf-java io.trino:trino-cache + + org.apache.arrow:arrow-memory-unsafe org.gaul:modernizer-maven-annotations diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/BlobUtils.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/BlobUtils.java index 57da99a..24bafb0 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/BlobUtils.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/BlobUtils.java @@ -32,7 +32,7 @@ public enum BlobVirtualColumnType { NONE, POSITION, - SIZE + SIZE, } private BlobUtils() {} diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceArrowToPageScanner.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceArrowToPageScanner.java index 47ab5d2..096bfdc 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceArrowToPageScanner.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceArrowToPageScanner.java @@ -158,53 +158,37 @@ public LanceArrowToPageScanner( } } - LanceScanner openedScanner = null; - ArrowReader openedReader = null; + LanceScanner openedScanner = scannerFactory.open( + path, + allocator, + projectionColumns, + storageOptions, + substraitFilter, + limit, + userIdentity, + datasetVersion); + ArrowReader openedReader; try { - openedScanner = scannerFactory.open(path, allocator, projectionColumns, storageOptions, substraitFilter, limit, userIdentity, datasetVersion); openedReader = openedScanner.scanBatches(); - VectorSchemaRoot openedRoot = openedReader.getVectorSchemaRoot(); - this.lanceScanner = openedScanner; - this.arrowReader = openedReader; - this.vectorSchemaRoot = openedRoot; } - catch (IOException e) { - RuntimeException failure = new RuntimeException("Unable to get vector schema root", e); - closeSuppressing(openedReader, failure); - closeFactorySuppressing(scannerFactory, failure); - throw failure; - } - catch (RuntimeException e) { - closeSuppressing(openedReader, e); - closeFactorySuppressing(scannerFactory, e); + catch (RuntimeException | Error e) { + closeScannerFactory(e, scannerFactory); throw e; } - } - - private static void closeSuppressing(AutoCloseable closeable, RuntimeException failure) - { - if (closeable == null) { - return; - } try { - closeable.close(); - } - catch (Exception closeFailure) { - failure.addSuppressed(closeFailure); - } - } - - private static void closeFactorySuppressing(ScannerFactory scannerFactory, RuntimeException failure) - { - if (scannerFactory == null) { - return; + this.vectorSchemaRoot = openedReader.getVectorSchemaRoot(); } - try { - scannerFactory.close(); - } - catch (RuntimeException closeFailure) { - failure.addSuppressed(closeFailure); + catch (Throwable t) { + closeArrowReader(t, openedReader); + closeScannerFactory(t, scannerFactory); + if (t instanceof IOException) { + throw new RuntimeException("Unable to get vector schema root", t); + } + throwIfUnchecked(t); + throw new RuntimeException(t); } + this.lanceScanner = openedScanner; + this.arrowReader = openedReader; } private long lastBatchBytes; @@ -231,8 +215,10 @@ public boolean read() catch (IllegalStateException e) { // Handle allocator closed during concurrent operations if (e.getMessage() != null && e.getMessage().contains("allocator")) { - throw new TrinoException(TRANSACTION_CONFLICT, - "Concurrent operation conflict: allocator was closed during read", e); + throw new TrinoException( + TRANSACTION_CONFLICT, + "Concurrent operation conflict: allocator was closed during read", + e); } throw e; } @@ -290,8 +276,12 @@ public void convert(PageBuilder pageBuilder) if (projectedVector == null) { throw new TrinoException(GENERIC_INTERNAL_ERROR, "Projected column not found in Lance scan: " + colHandle.path()); } - convertType(pageBuilder.getBlockBuilder(column), columnTypes.get(column), projectedVector.vector(), 0, - rowCount, projectedVector.nullChecker()); + convertType(pageBuilder.getBlockBuilder(column), + columnTypes.get(column), + projectedVector.vector(), + 0, + rowCount, + projectedVector.nullChecker()); } } vectorSchemaRoot.clear(); @@ -401,53 +391,84 @@ private void convertType(BlockBuilder output, Type type, FieldVector vector, int Class javaType = type.getJavaType(); try { if (javaType == boolean.class) { - writeVectorValues(output, nullChecker, - index -> type.writeBoolean(output, ((BitVector) vector).get(index) == 1), offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeBoolean(output, ((BitVector) vector).get(index) == 1), + offset, + length); } else if (javaType == long.class) { if (type.equals(BIGINT)) { // Handle both signed (BigIntVector) and unsigned (UInt8Vector) 64-bit integers, // and unsigned 32-bit integers (UInt4Vector) promoted to BIGINT if (vector instanceof UInt8Vector uint8Vector) { - writeVectorValues(output, nullChecker, - index -> type.writeLong(output, uint8Vector.get(index)), offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, uint8Vector.get(index)), + offset, + length); } else if (vector instanceof UInt4Vector uint4Vector) { - writeVectorValues(output, nullChecker, - index -> type.writeLong(output, Integer.toUnsignedLong(uint4Vector.get(index))), offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, Integer.toUnsignedLong(uint4Vector.get(index))), + offset, + length); } else { - writeVectorValues(output, nullChecker, - index -> type.writeLong(output, ((BigIntVector) vector).get(index)), offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, ((BigIntVector) vector).get(index)), + offset, + length); } } else if (type.equals(INTEGER)) { if (vector instanceof UInt4Vector uint4Vector) { - writeVectorValues(output, nullChecker, - index -> type.writeLong(output, Integer.toUnsignedLong(uint4Vector.get(index))), offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, Integer.toUnsignedLong(uint4Vector.get(index))), + offset, + length); } else { - writeVectorValues(output, nullChecker, index -> type.writeLong(output, ((IntVector) vector).get(index)), - offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, ((IntVector) vector).get(index)), + offset, + length); } } else if (type.equals(DATE)) { - writeVectorValues(output, nullChecker, - index -> type.writeLong(output, ((DateDayVector) vector).get(index)), offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, ((DateDayVector) vector).get(index)), + offset, + length); } else if (type.equals(TIME_MICROS)) { - writeVectorValues(output, nullChecker, index -> type.writeLong(output, + writeVectorValues(output, nullChecker, index -> type.writeLong( + output, ((TimeMicroVector) vector).get(index) * PICOSECONDS_PER_MICROSECOND), offset, length); } else if (type.equals(REAL)) { // REAL stores float bits as int which is widened to long if (vector instanceof Float2Vector f2v) { // Widen float16 to float32 since Trino has no float16 type - writeVectorValues(output, nullChecker, index -> type.writeLong(output, + writeVectorValues(output, nullChecker, index -> type.writeLong( + output, Float.floatToIntBits(f2v.getValueAsFloat(index))), offset, length); } else { - writeVectorValues(output, nullChecker, index -> type.writeLong(output, + writeVectorValues(output, nullChecker, index -> type.writeLong( + output, Float.floatToIntBits(((Float4Vector) vector).get(index))), offset, length); } } @@ -470,34 +491,49 @@ else if (vector instanceof TimeStampMicroVector tsVector) { }, offset, length); } else { - throw new TrinoException(GENERIC_INTERNAL_ERROR, + throw new TrinoException( + GENERIC_INTERNAL_ERROR, format("Expected TimeStampMicroTZVector but got: %s", vector.getClass().getSimpleName())); } } else if (type instanceof TimestampType) { // Timestamp without timezone - stored as microseconds in Arrow if (vector instanceof TimeStampMicroVector tsVector) { - writeVectorValues(output, nullChecker, index -> type.writeLong(output, tsVector.get(index)), - offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, tsVector.get(index)), + offset, + length); } else if (vector instanceof TimeStampMicroTZVector tsVector) { // Handle case where Arrow has TZ but Trino doesn't need it - writeVectorValues(output, nullChecker, index -> type.writeLong(output, tsVector.get(index)), - offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeLong(output, tsVector.get(index)), + offset, + length); } else { - throw new TrinoException(GENERIC_INTERNAL_ERROR, + throw new TrinoException( + GENERIC_INTERNAL_ERROR, format("Expected TimeStampMicroVector but got: %s", vector.getClass().getSimpleName())); } } else { - throw new TrinoException(GENERIC_INTERNAL_ERROR, + throw new TrinoException( + GENERIC_INTERNAL_ERROR, format("Unhandled type for %s: %s", javaType.getSimpleName(), type)); } } else if (javaType == double.class) { - writeVectorValues(output, nullChecker, index -> type.writeDouble(output, ((Float8Vector) vector).get(index)), - offset, length); + writeVectorValues( + output, + nullChecker, + index -> type.writeDouble(output, ((Float8Vector) vector).get(index)), + offset, + length); } else if (javaType == Slice.class) { writeVectorValues(output, nullChecker, index -> writeSlice(output, type, vector, index), offset, length); @@ -505,36 +541,59 @@ else if (javaType == Slice.class) { else if (type instanceof ArrayType arrayType) { // Handle both ListVector and FixedSizeListVector if (vector instanceof FixedSizeListVector) { - writeVectorValues(output, nullChecker, index -> writeFixedSizeArrayBlock(output, arrayType, vector, index), offset, + writeVectorValues( + output, + nullChecker, + index -> writeFixedSizeArrayBlock(output, arrayType, vector, index), + offset, length); } else { - writeVectorValues(output, nullChecker, index -> writeArrayBlock(output, arrayType, vector, index), offset, + writeVectorValues( + output, + nullChecker, + index -> writeArrayBlock(output, arrayType, vector, index), + offset, length); } } else if (type instanceof RowType rowType) { - writeVectorValues(output, nullChecker, index -> writeRowBlock(output, rowType, vector, index), offset, + writeVectorValues( + output, + nullChecker, + index -> writeRowBlock(output, rowType, vector, index), + offset, length); } else { - throw new TrinoException(GENERIC_INTERNAL_ERROR, + throw new TrinoException( + GENERIC_INTERNAL_ERROR, format("Unhandled type for %s: %s", javaType.getSimpleName(), type)); } } catch (ClassCastException ex) { - throw new TrinoException(GENERIC_INTERNAL_ERROR, - format("Unhandled type for %s: %s", javaType.getSimpleName(), type), ex); + throw new TrinoException( + GENERIC_INTERNAL_ERROR, + format("Unhandled type for %s: %s", javaType.getSimpleName(), type), + ex); } } - private void writeVectorValues(BlockBuilder output, FieldVector vector, Consumer consumer, int offset, + private void writeVectorValues( + BlockBuilder output, + FieldVector vector, + Consumer consumer, + int offset, int length) { writeVectorValues(output, vector::isNull, consumer, offset, length); } - private void writeVectorValues(BlockBuilder output, IntPredicate nullChecker, Consumer consumer, int offset, + private void writeVectorValues( + BlockBuilder output, + IntPredicate nullChecker, + Consumer consumer, + int offset, int length) { for (int i = offset; i < offset + length; i++) { @@ -640,13 +699,62 @@ private void writeRowBlock(BlockBuilder output, RowType rowType, FieldVector vec @Override public void close() { - vectorSchemaRoot.close(); + Throwable failure = null; + try { + vectorSchemaRoot.close(); + } + catch (Throwable t) { + failure = t; + } + + failure = closeArrowReader(failure, arrowReader); + failure = closeScannerFactory(failure, scannerFactory); + + if (failure != null) { + throwIfUnchecked(failure); + throw new RuntimeException(failure); + } + } + + private static Throwable closeArrowReader(Throwable failure, ArrowReader arrowReader) + { try { arrowReader.close(); } - catch (IOException ioe) { - // ignore for now. + catch (Throwable t) { + if (failure == null) { + failure = t; + } + else { + failure.addSuppressed(t); + } + } + return failure; + } + + private static Throwable closeScannerFactory(Throwable failure, ScannerFactory scannerFactory) + { + try { + scannerFactory.close(); + } + catch (Throwable t) { + if (failure == null) { + failure = t; + } + else { + failure.addSuppressed(t); + } + } + return failure; + } + + private static void throwIfUnchecked(Throwable failure) + { + if (failure instanceof RuntimeException runtimeException) { + throw runtimeException; + } + if (failure instanceof Error error) { + throw error; } - scannerFactory.close(); } } diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceBasePageSource.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceBasePageSource.java index 5bb493c..46d59fa 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceBasePageSource.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceBasePageSource.java @@ -17,6 +17,7 @@ import io.trino.spi.PageBuilder; import io.trino.spi.TrinoException; import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.connector.SourcePage; import org.apache.arrow.memory.BufferAllocator; import java.util.List; @@ -87,8 +88,8 @@ private static boolean isConcurrentModificationError(RuntimeException e) String message = current.getMessage(); if (message != null && ( message.toLowerCase().contains("not found") || - message.toLowerCase().contains("concurrent") || - message.toLowerCase().contains("conflict"))) { + message.toLowerCase().contains("concurrent") || + message.toLowerCase().contains("conflict"))) { return true; } if (current instanceof NullPointerException) { @@ -118,7 +119,7 @@ public boolean isFinished() } @Override - public Page getNextPage() + public SourcePage getNextSourcePage() { checkState(pageBuilder.isEmpty(), "PageBuilder is not empty at the beginning of a new page"); if (!lanceArrowToPageScanner.read()) { @@ -130,7 +131,7 @@ public Page getNextPage() lanceArrowToPageScanner.convert(pageBuilder); Page page = pageBuilder.build(); pageBuilder.reset(); - return page; + return SourcePage.create(page); } @Override @@ -142,8 +143,10 @@ public long getMemoryUsage() @Override public void close() { - lanceArrowToPageScanner.close(); - // Close the child allocator - this releases resources allocated by this page source - bufferAllocator.close(); + // Close the child allocator even if the scanner fails to close, so the resources allocated + // by this page source are always released (scanner close() can throw). + try (bufferAllocator) { + lanceArrowToPageScanner.close(); + } } } diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceColumnHandle.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceColumnHandle.java index f370090..afa9c1d 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceColumnHandle.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceColumnHandle.java @@ -92,8 +92,7 @@ public LanceColumnHandle( BlobUtils.BlobVirtualColumnType blobVirtualColumnType, String baseBlobColumnName) { - this( - name, + this(name, trinoType, isNullable, fieldId, diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnector.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnector.java index 5270307..ff8d7ad 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnector.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnector.java @@ -54,7 +54,9 @@ public LanceConnector( } @Override - public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel, boolean readOnly, + public ConnectorTransactionHandle beginTransaction( + IsolationLevel isolationLevel, + boolean readOnly, boolean autoCommit) { return LanceTransactionHandle.INSTANCE; diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnectorFactory.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnectorFactory.java index 1afcd77..1d5b1e2 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnectorFactory.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceConnectorFactory.java @@ -21,6 +21,7 @@ import com.google.inject.TypeLiteral; import io.airlift.bootstrap.Bootstrap; import io.airlift.json.JsonModule; +import io.trino.plugin.base.ConnectorContextModule; import io.trino.plugin.base.TypeDeserializerModule; import io.trino.plugin.base.jmx.MBeanServerModule; import io.trino.spi.connector.Connector; @@ -82,8 +83,9 @@ public Connector create(String catalogName, Map config, Connecto ImmutableList.Builder modulesBuilder = ImmutableList.builder().add(new JsonModule()).add(new MBeanModule()) - .add(new MBeanServerModule()).add(new TypeDeserializerModule(context.getTypeManager())) + .add(new MBeanServerModule()).add(new TypeDeserializerModule()) .add(new LanceModule()) + .add(new ConnectorContextModule(catalogName, context)) // Bind the raw namespace properties map for free-form property access .add(binder -> binder.bind(new TypeLiteral>() {}) .annotatedWith(LanceNamespaceProperties.class) diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCountPageSource.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCountPageSource.java index 3f18397..51e4f92 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCountPageSource.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCountPageSource.java @@ -17,6 +17,7 @@ import io.trino.spi.Page; import io.trino.spi.block.BlockBuilder; import io.trino.spi.connector.ConnectorPageSource; +import io.trino.spi.connector.SourcePage; import org.lance.ManifestSummary; import java.util.Map; @@ -72,7 +73,7 @@ public boolean isFinished() } @Override - public Page getNextPage() + public SourcePage getNextSourcePage() { if (finished.get()) { return null; @@ -84,7 +85,7 @@ public Page getNextPage() BlockBuilder blockBuilder = BIGINT.createBlockBuilder(null, 1); BIGINT.writeLong(blockBuilder, count); - return new Page(blockBuilder.build()); + return SourcePage.create(new Page(blockBuilder.build())); } private long computeCount() diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceFragmentPageSource.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceFragmentPageSource.java index d10b3d3..eba0829 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceFragmentPageSource.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceFragmentPageSource.java @@ -15,7 +15,6 @@ import io.airlift.log.Logger; import org.apache.arrow.memory.BufferAllocator; -import org.lance.Dataset; import org.lance.ipc.LanceScanner; import org.lance.ipc.ScanOptions; @@ -97,7 +96,7 @@ public static class FragmentScannerFactory private final boolean includeRowAddress; private final int readBatchSize; private final LanceRuntime runtime; - private Dataset lanceDataset; + private LanceRuntime.DatasetLease datasetLease; private LanceScanner lanceScanner; public FragmentScannerFactory(List fragmentIds, boolean includeRowAddress, int readBatchSize, LanceRuntime runtime) @@ -114,9 +113,15 @@ public FragmentScannerFactory(Optional> fragmentIds, boolean inclu } @Override - public LanceScanner open(String tablePath, BufferAllocator allocator, List columns, - Map storageOptions, Optional substraitFilter, OptionalLong limit, - String userIdentity, Long datasetVersion) + public LanceScanner open( + String tablePath, + BufferAllocator allocator, + List columns, + Map storageOptions, + Optional substraitFilter, + OptionalLong limit, + String userIdentity, + Long datasetVersion) { ScanOptions.Builder optionsBuilder = new ScanOptions.Builder(); if (!columns.isEmpty()) { @@ -132,7 +137,8 @@ public LanceScanner open(String tablePath, BufferAllocator allocator, List Integer.toString(ids.size())).orElse("all"), readBatchSize, substraitFilter.isPresent() ? "present" : "none", @@ -141,14 +147,18 @@ public LanceScanner open(String tablePath, BufferAllocator allocator, List new ArrayList<>()) + deletedRows.computeIfAbsent(fragmentId, _ -> new ArrayList<>()) .add(rowIndex); deletePositions.add(position); rowsDeleted++; @@ -139,8 +139,11 @@ public void storeMergedRows(Page page) insertPageSink.appendPage(insertPage); } - log.debug("storeMergedRows: processed %d rows, deleted=%d, inserted=%d", - page.getPositionCount(), deletePositions.size(), insertPositions.size()); + log.debug( + "storeMergedRows: processed %d rows, deleted=%d, inserted=%d", + page.getPositionCount(), + deletePositions.size(), + insertPositions.size()); } /** @@ -161,8 +164,11 @@ private Page extractDataPage(Page page, List positions, int dataColumnC @Override public CompletableFuture> finish() { - log.debug("finish: completing merge with %d deletions across %d fragments, %d inserts", - rowsDeleted, deletedRows.size(), rowsInserted); + log.debug( + "finish: completing merge with %d deletions across %d fragments, %d inserts", + rowsDeleted, + deletedRows.size(), + rowsInserted); Collection insertResults = insertPageSink.finish().join(); diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceMetadata.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceMetadata.java index e4a1af7..2ca665a 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceMetadata.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceMetadata.java @@ -67,7 +67,6 @@ import org.lance.Dataset; import org.lance.FragmentMetadata; import org.lance.ManifestSummary; -import org.lance.ReadOptions; import org.lance.SourcedTransaction; import org.lance.Transaction; import org.lance.namespace.LanceNamespace; @@ -109,6 +108,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.UUID; @@ -137,12 +137,12 @@ public class LanceMetadata implements ConnectorMetadata { private static final Logger log = Logger.get(LanceMetadata.class); - private static final ConcurrentMap transactionDatasets = new ConcurrentHashMap<>(); private final LanceRuntime runtime; private final LanceConfig lanceConfig; private final JsonCodec commitTaskDataCodec; private final JsonCodec mergeCommitDataCodec; + private final ConcurrentMap transactionDatasets = new ConcurrentHashMap<>(); @Inject public LanceMetadata( @@ -157,6 +157,18 @@ public LanceMetadata( this.mergeCommitDataCodec = requireNonNull(mergeCommitDataCodec, "mergeCommitDataCodec is null"); } + @Override + public void cleanupQuery(ConnectorSession session) + { + String queryId = session.getQueryId(); + transactionDatasets.forEach((transactionId, transactionDataset) -> { + if (Objects.equals(transactionDataset.queryId(), queryId) && + transactionDatasets.remove(transactionId, transactionDataset)) { + closeTransactionDataset(transactionId, transactionDataset); + } + }); + } + // ===== Schema/Namespace Operations ===== @Override @@ -168,7 +180,7 @@ public List listSchemaNames(ConnectorSession session) ListNamespacesRequest request = new ListNamespacesRequest(); if (runtime.getParentPrefix().isPresent()) { - request.setId(runtime.getParentPrefix().get()); + request.setId(runtime.getParentPrefix().orElseThrow()); } ListNamespacesResponse response = getNamespace().listNamespaces(request); Set namespaces = response.getNamespaces(); @@ -247,8 +259,11 @@ public Map getSchemaProperties(ConnectorSession session, String // ===== Table Operations ===== @Override - public LanceTableHandle getTableHandle(ConnectorSession session, SchemaTableName name, - Optional startVersion, Optional endVersion) + public LanceTableHandle getTableHandle( + ConnectorSession session, + SchemaTableName name, + Optional startVersion, + Optional endVersion) { if (startVersion.isPresent()) { throw new TrinoException(NOT_SUPPORTED, "Lance connector does not support start version for time travel"); @@ -262,7 +277,7 @@ public LanceTableHandle getTableHandle(ConnectorSession session, SchemaTableName Long datasetVersion; if (endVersion.isPresent()) { - datasetVersion = resolveVersion(session, tablePath, storageOptions, endVersion.get()); + datasetVersion = resolveVersion(session, tablePath, storageOptions, endVersion.orElseThrow()); } else { datasetVersion = runtime.getLatestVersion(userIdentity, tablePath, storageOptions); @@ -273,8 +288,11 @@ public LanceTableHandle getTableHandle(ConnectorSession session, SchemaTableName return null; } - private Long resolveVersion(ConnectorSession session, String tablePath, - Map storageOptions, ConnectorTableVersion version) + private Long resolveVersion( + ConnectorSession session, + String tablePath, + Map storageOptions, + ConnectorTableVersion version) { String userIdentity = session.getUser(); Type versionType = version.getVersionType(); @@ -285,8 +303,12 @@ private Long resolveVersion(ConnectorSession session, String tablePath, }; } - private Long resolveTargetIdVersion(String tablePath, Map storageOptions, - String userIdentity, ConnectorTableVersion version, Type versionType) + private Long resolveTargetIdVersion( + String tablePath, + Map storageOptions, + String userIdentity, + ConnectorTableVersion version, + Type versionType) { long versionNumber; if (versionType.equals(BIGINT)) { @@ -318,8 +340,13 @@ else if (versionType.equals(TINYINT)) { return versionNumber; } - private Long resolveTemporalVersion(ConnectorSession session, String tablePath, - Map storageOptions, String userIdentity, ConnectorTableVersion version, Type versionType) + private Long resolveTemporalVersion( + ConnectorSession session, + String tablePath, + Map storageOptions, + String userIdentity, + ConnectorTableVersion version, + Type versionType) { long timestampMillis = getTimestampMillis(session, version, versionType); @@ -327,13 +354,16 @@ private Long resolveTemporalVersion(ConnectorSession session, String tablePath, userIdentity, tablePath, timestampMillis, storageOptions); if (resolvedVersion.isEmpty()) { - throw new TrinoException(INVALID_ARGUMENTS, + throw new TrinoException( + INVALID_ARGUMENTS, "No Lance version found at or before timestamp: " + Instant.ofEpochMilli(timestampMillis)); } - log.debug("Resolved temporal version for timestamp %s to version %d", - Instant.ofEpochMilli(timestampMillis), resolvedVersion.get()); - return resolvedVersion.get(); + log.debug( + "Resolved temporal version for timestamp %s to version %d", + Instant.ofEpochMilli(timestampMillis), + resolvedVersion.orElseThrow()); + return resolvedVersion.orElseThrow(); } private long getTimestampMillis(ConnectorSession session, ConnectorTableVersion version, Type versionType) @@ -365,7 +395,8 @@ private long getTimestampMillis(ConnectorSession session, ConnectorTableVersion : ((LongTimestampWithTimeZone) version.getVersion()).getEpochMillis(); } - throw new TrinoException(NOT_SUPPORTED, + throw new TrinoException( + NOT_SUPPORTED, "Unsupported type for Lance temporal version: " + versionType.getDisplayName()); } @@ -418,8 +449,11 @@ public Map getColumnHandles(ConnectorSession session, Conn try { Map storageOptions = getEffectiveStorageOptions(lanceTableHandle); String userIdentity = session.getUser(); - return runtime.getColumnHandles(userIdentity, lanceTableHandle.getTablePath(), - lanceTableHandle.getDatasetVersion(), storageOptions); + return runtime.getColumnHandles( + userIdentity, + lanceTableHandle.getTablePath(), + lanceTableHandle.getDatasetVersion(), + storageOptions); } catch (Exception e) { throw new TableNotFoundException(new SchemaTableName(lanceTableHandle.getSchemaName(), lanceTableHandle.getTableName())); @@ -427,7 +461,8 @@ public Map getColumnHandles(ConnectorSession session, Conn } @Override - public Map> listTableColumns(ConnectorSession session, + public Map> listTableColumns( + ConnectorSession session, SchemaTablePrefix prefix) { requireNonNull(prefix, "prefix is null"); @@ -454,15 +489,20 @@ public Map> listTableColumns(ConnectorSess } @Override - public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTableHandle tableHandle, + public ColumnMetadata getColumnMetadata( + ConnectorSession session, + ConnectorTableHandle tableHandle, ColumnHandle columnHandle) { return ((LanceColumnHandle) columnHandle).getColumnMetadata(); } @Override - public Optional> applyProjection(ConnectorSession session, - ConnectorTableHandle handle, List projections, Map assignments) + public Optional> applyProjection( + ConnectorSession session, + ConnectorTableHandle handle, + List projections, + Map assignments) { LanceTableHandle lanceTableHandle = (LanceTableHandle) handle; List projectedExpressions = new ArrayList<>(); @@ -475,7 +515,7 @@ public Optional> applyProjecti return Optional.empty(); } - LanceColumnHandle lanceColumn = column.get(); + LanceColumnHandle lanceColumn = column.orElseThrow(); String variableName = lanceColumn.path(); projectedColumns.putIfAbsent(variableName, lanceColumn); projectedExpressions.add(new Variable(variableName, lanceColumn.trinoType())); @@ -515,7 +555,7 @@ private static Optional resolveProjectionColumn( if (target.isEmpty()) { return Optional.empty(); } - LanceColumnHandle targetColumn = target.get(); + LanceColumnHandle targetColumn = target.orElseThrow(); if (!(fieldDereference.getTarget().getType() instanceof RowType rowType)) { return Optional.empty(); } @@ -533,7 +573,7 @@ private static Optional resolveProjectionColumn( List dereferencePath = new ArrayList<>(targetColumn.dereferencePath()); dereferencePath.add(fieldIndex); List dereferenceNames = new ArrayList<>(targetColumn.dereferenceNames()); - dereferenceNames.add(field.getName().get()); + dereferenceNames.add(field.getName().orElseThrow()); String canonicalPath = LanceFieldPath.canonicalPath(buildFieldPath(targetColumn.baseColumnName(), dereferenceNames)); return Optional.of(LanceColumnHandle.nestedColumn( @@ -569,7 +609,8 @@ public TableStatistics getTableStatistics(ConnectorSession session, ConnectorTab ManifestSummary summary = runtime.getManifestSummary( userIdentity, lanceTableHandle.getTablePath(), lanceTableHandle.getDatasetVersion(), storageOptions); - log.debug("getTableStatistics: table=%s, totalRows=%d, totalFilesSize=%d, totalFragments=%d", + log.debug( + "getTableStatistics: table=%s, totalRows=%d, totalFilesSize=%d, totalFragments=%d", lanceTableHandle.getTableName(), summary.getTotalRows(), summary.getTotalFilesSize(), @@ -588,8 +629,10 @@ public TableStatistics getTableStatistics(ConnectorSession session, ConnectorTab } @Override - public Optional> applyLimit(ConnectorSession session, - ConnectorTableHandle table, long limit) + public Optional> applyLimit( + ConnectorSession session, + ConnectorTableHandle table, + long limit) { LanceTableHandle lanceTableHandle = (LanceTableHandle) table; @@ -640,7 +683,8 @@ public Optional> applyAggrega return Optional.empty(); } - log.debug("applyAggregation: pushing COUNT(*) for table %s (no filter)", + log.debug( + "applyAggregation: pushing COUNT(*) for table %s (no filter)", lanceTableHandle.getTableName()); LanceTableHandle newHandle = lanceTableHandle.withCountStar(); @@ -665,8 +709,10 @@ private boolean isCountStar(AggregateFunction aggregate) } @Override - public Optional> applyFilter(ConnectorSession session, - ConnectorTableHandle table, Constraint constraint) + public Optional> applyFilter( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint) { LanceTableHandle lanceTableHandle = (LanceTableHandle) table; @@ -691,8 +737,12 @@ public Optional> applyFilter(C List exprColumnNames = exprResult.columnNames(); io.trino.spi.expression.ConnectorExpression remainingExpression = exprResult.remainingExpression(); - log.debug("applyFilter: newConstraint=%s, pushedTupleDomain=%s, remainingTupleDomain=%s, pushedExpressions=%d", - newConstraint, tupleDomainResult.pushedTupleDomain(), tupleDomainResult.remainingTupleDomain(), pushedExpressions.size()); + log.debug( + "applyFilter: newConstraint=%s, pushedTupleDomain=%s, remainingTupleDomain=%s, pushedExpressions=%d", + newConstraint, + tupleDomainResult.pushedTupleDomain(), + tupleDomainResult.remainingTupleDomain(), + pushedExpressions.size()); // If no TupleDomain constraints and no pushed expressions, nothing to push down if (tupleDomainResult.expression().isEmpty() && pushedExpressions.isEmpty()) { @@ -710,7 +760,7 @@ public Optional> applyFilter(C } // Combine with existing filter if present - byte[] newFilterBytes = substraitFilter.get().array(); + byte[] newFilterBytes = substraitFilter.orElseThrow().array(); byte[] existingFilter = lanceTableHandle.getSubstraitFilter(); // If there's an existing filter, we can't easily combine Substrait expressions, @@ -736,8 +786,11 @@ public Optional> applyFilter(C LanceTableHandle newHandle = lanceTableHandle.withSubstraitFilter(newFilterBytes, filterColumnNames); TupleDomain remainingFilter = tupleDomainResult.remainingTupleDomain(); - log.debug("applyFilter: pushing substrait filter (size=%d bytes, %d expressions), remaining=%s", - newFilterBytes.length, pushedExpressions.size(), remainingFilter); + log.debug( + "applyFilter: pushing substrait filter (size=%d bytes, %d expressions), remaining=%s", + newFilterBytes.length, + pushedExpressions.size(), + remainingFilter); return Optional.of(new ConstraintApplicationResult<>( newHandle, @@ -861,8 +914,8 @@ else if (saveMode == SaveMode.IGNORE) { Schema arrowSchema = LancePageToArrowConverter.toArrowSchema(tableMetadata.getColumns(), blobColumns, vectorColumns); String userIdentity = session.getUser(); // For write operations, open dataset directly (not cached) - try (Dataset dataset = runtime.openDatasetDirect(userIdentity, existingPath, null, storageOptions)) { - commitOverwrite(dataset, List.of(), arrowSchema, storageOptions); + try (LanceRuntime.DatasetLease datasetLease = runtime.openDatasetDirectLease(userIdentity, existingPath, null, storageOptions)) { + commitOverwrite(datasetLease.getDataset(), List.of(), arrowSchema, storageOptions); } runtime.invalidate(userIdentity, existingPath); log.debug("createTable: replaced table %s at %s", tableName, existingPath); @@ -938,8 +991,12 @@ public ConnectorOutputTableHandle beginCreateTable( } List columns = tableMetadata.getColumns().stream() - .map(col -> new LanceColumnHandle(col.getName(), col.getType(), col.isNullable(), - -1, blobColumns.contains(col.getName()))) + .map(col -> new LanceColumnHandle( + col.getName(), + col.getType(), + col.isNullable(), + -1, + blobColumns.contains(col.getName()))) .collect(toImmutableList()); Schema arrowSchema = LancePageToArrowConverter.toArrowSchema(tableMetadata.getColumns(), blobColumns, vectorColumns); @@ -948,19 +1005,40 @@ public ConnectorOutputTableHandle beginCreateTable( String transactionId = null; if (tableExisted) { transactionId = UUID.randomUUID().toString(); - ReadOptions readOptions = new ReadOptions.Builder() - .setStorageOptions(storageOptions) - .build(); - Dataset dataset = Dataset.open(tablePath, readOptions); - transactionDatasets.put(transactionId, dataset); - // For replace tables (RTAS/CORTAS), use existing table's format if not specified - if (fileFormatVersion == null) { - fileFormatVersion = dataset.getLanceFileFormatVersion(); + LanceRuntime.DatasetLease datasetLease = runtime.openDatasetDirectLease(session.getUser(), tablePath, null, storageOptions); + boolean registered = false; + try { + registerTransactionDataset(transactionId, session, datasetLease); + registered = true; + // For replace tables (RTAS/CORTAS), use existing table's format if not specified + if (fileFormatVersion == null) { + fileFormatVersion = datasetLease.getDataset().getLanceFileFormatVersion(); + } + datasetLease = null; + } + catch (RuntimeException | Error e) { + if (registered) { + TransactionDataset transactionDataset = transactionDatasets.remove(transactionId); + if (transactionDataset != null) { + closeTransactionDataset(transactionId, transactionDataset); + } + } + else { + datasetLease.close(); + } + throw e; } } - log.debug("beginCreateTable: table=%s, path=%s, replace=%s, tableExisted=%s, transactionId=%s, blobColumns=%s, fileFormatVersion=%s", - tableName, tablePath, replace, tableExisted, transactionId, blobColumns, fileFormatVersion); + log.debug( + "beginCreateTable: table=%s, path=%s, replace=%s, tableExisted=%s, transactionId=%s, blobColumns=%s, fileFormatVersion=%s", + tableName, + tablePath, + replace, + tableExisted, + transactionId, + blobColumns, + fileFormatVersion); return new LanceWritableTableHandle( tableName, @@ -993,16 +1071,19 @@ public Optional finishCreateTable( } String transactionId = handle.transactionId(); - log.debug("finishCreateTable: table=%s, fragments=%d, replace=%s, tableExisted=%s, transactionId=%s", - handle.tableName(), fragments.size(), handle.replace(), handle.tableExisted(), transactionId); + log.debug( + "finishCreateTable: table=%s, fragments=%d, replace=%s, tableExisted=%s, transactionId=%s", + handle.tableName(), + fragments.size(), + handle.replace(), + handle.tableExisted(), + transactionId); Map storageOptions = handle.storageOptions(); if (handle.tableExisted()) { - Dataset dataset = transactionDatasets.remove(transactionId); - if (dataset == null) { - throw new TrinoException(GENERIC_INTERNAL_ERROR, "No dataset found for transaction: " + transactionId); - } + TransactionDataset transactionDataset = removeTransactionDataset(transactionId); + Dataset dataset = transactionDataset.dataset(); try { if (fragments.isEmpty()) { commitOverwrite(dataset, List.of(), arrowSchema, storageOptions); @@ -1013,7 +1094,7 @@ public Optional finishCreateTable( } } finally { - dataset.close(); + closeTransactionDataset(transactionId, transactionDataset); } } else { @@ -1053,26 +1134,48 @@ public ConnectorInsertTableHandle beginInsert( String transactionId = UUID.randomUUID().toString(); // For write operations, open dataset directly (not cached) - Dataset dataset = runtime.openDatasetDirect(userIdentity, tablePath, null, storageOptions); - transactionDatasets.put(transactionId, dataset); - - // Read the existing table's file format version to ensure consistent writes - String fileFormatVersion = dataset.getLanceFileFormatVersion(); - log.debug("beginInsert: table=%s, path=%s, columns=%d, transactionId=%s, fileFormatVersion=%s", - tableName, tablePath, columns.size(), transactionId, fileFormatVersion); - - return new LanceWritableTableHandle( - tableName, - tablePath, - schemaJson, - lanceColumns, - tableId, - storageOptions, - false, - false, - true, - transactionId, - fileFormatVersion); + LanceRuntime.DatasetLease datasetLease = runtime.openDatasetDirectLease(userIdentity, tablePath, null, storageOptions); + boolean registered = false; + try { + registerTransactionDataset(transactionId, session, datasetLease); + registered = true; + + // Read the existing table's file format version to ensure consistent writes + String fileFormatVersion = datasetLease.getDataset().getLanceFileFormatVersion(); + log.debug( + "beginInsert: table=%s, path=%s, columns=%d, transactionId=%s, fileFormatVersion=%s", + tableName, + tablePath, + columns.size(), + transactionId, + fileFormatVersion); + + datasetLease = null; + return new LanceWritableTableHandle( + tableName, + tablePath, + schemaJson, + lanceColumns, + tableId, + storageOptions, + false, + false, + true, + transactionId, + fileFormatVersion); + } + catch (RuntimeException | Error e) { + if (registered) { + TransactionDataset transactionDataset = transactionDatasets.remove(transactionId); + if (transactionDataset != null) { + closeTransactionDataset(transactionId, transactionDataset); + } + } + else { + datasetLease.close(); + } + throw e; + } } @Override @@ -1088,10 +1191,8 @@ public Optional finishInsert( log.debug("finishInsert: table=%s, fragments=%d, transactionId=%s", handle.tableName(), fragments.size(), transactionId); - Dataset dataset = transactionDatasets.remove(transactionId); - if (dataset == null) { - throw new TrinoException(GENERIC_INTERNAL_ERROR, "No dataset found for transaction: " + transactionId); - } + TransactionDataset transactionDataset = removeTransactionDataset(transactionId); + Dataset dataset = transactionDataset.dataset(); try { if (fragments.isEmpty()) { @@ -1107,7 +1208,7 @@ public Optional finishInsert( return Optional.empty(); } finally { - dataset.close(); + closeTransactionDataset(transactionId, transactionDataset); } } @@ -1149,27 +1250,50 @@ public ConnectorMergeTableHandle beginMerge( String transactionId = UUID.randomUUID().toString(); String userIdentity = session.getUser(); // For write operations, open dataset directly (not cached) - Dataset dataset = runtime.openDatasetDirect(userIdentity, tablePath, null, storageOptions); - long readVersion = dataset.version(); - transactionDatasets.put(transactionId, dataset); - - List columns = runtime.getColumnHandleList(userIdentity, tablePath, null, storageOptions); - Schema arrowSchema = runtime.getSchema(userIdentity, tablePath, null, storageOptions); - String schemaJson = arrowSchema.toJson(); - - // Read the existing table's file format version to ensure consistent writes - String fileFormatVersion = dataset.getLanceFileFormatVersion(); - log.debug("beginMerge: table=%s, path=%s, version=%d, transactionId=%s, fileFormatVersion=%s", - tableName, tablePath, readVersion, transactionId, fileFormatVersion); - - return new LanceMergeTableHandle( - table.withStorageOptions(storageOptions), - getMergeRowIdColumnHandle(session, tableHandle), - readVersion, - schemaJson, - columns, - transactionId, - fileFormatVersion); + LanceRuntime.DatasetLease datasetLease = runtime.openDatasetDirectLease(userIdentity, tablePath, null, storageOptions); + boolean registered = false; + try { + Dataset dataset = datasetLease.getDataset(); + long readVersion = dataset.version(); + registerTransactionDataset(transactionId, session, datasetLease); + registered = true; + + List columns = runtime.getColumnHandleList(userIdentity, tablePath, null, storageOptions); + Schema arrowSchema = runtime.getSchema(userIdentity, tablePath, null, storageOptions); + String schemaJson = arrowSchema.toJson(); + + // Read the existing table's file format version to ensure consistent writes + String fileFormatVersion = dataset.getLanceFileFormatVersion(); + log.debug( + "beginMerge: table=%s, path=%s, version=%d, transactionId=%s, fileFormatVersion=%s", + tableName, + tablePath, + readVersion, + transactionId, + fileFormatVersion); + + datasetLease = null; + return new LanceMergeTableHandle( + table.withStorageOptions(storageOptions), + getMergeRowIdColumnHandle(session, tableHandle), + readVersion, + schemaJson, + columns, + transactionId, + fileFormatVersion); + } + catch (RuntimeException | Error e) { + if (registered) { + TransactionDataset transactionDataset = transactionDatasets.remove(transactionId); + if (transactionDataset != null) { + closeTransactionDataset(transactionId, transactionDataset); + } + } + else { + datasetLease.close(); + } + throw e; + } } @Override @@ -1183,13 +1307,14 @@ public void finishMerge( LanceMergeTableHandle handle = (LanceMergeTableHandle) mergeTableHandle; String transactionId = handle.transactionId(); - log.debug("finishMerge: table=%s, fragments=%d, transactionId=%s", - handle.tableHandle().getTableName(), fragments.size(), transactionId); + log.debug( + "finishMerge: table=%s, fragments=%d, transactionId=%s", + handle.tableHandle().getTableName(), + fragments.size(), + transactionId); - Dataset dataset = transactionDatasets.remove(transactionId); - if (dataset == null) { - throw new TrinoException(GENERIC_INTERNAL_ERROR, "No dataset found for transaction: " + transactionId); - } + TransactionDataset transactionDataset = removeTransactionDataset(transactionId); + Dataset dataset = transactionDataset.dataset(); try { List removedFragmentIds = new ArrayList<>(); @@ -1204,7 +1329,7 @@ public void finishMerge( LanceMergeCommitData commitData = mergeCommitDataCodec.fromJson(slice.getBytes()); for (FragmentDeletion deletion : commitData.deletions()) { - allDeletions.computeIfAbsent(deletion.fragmentId(), k -> new ArrayList<>()) + allDeletions.computeIfAbsent(deletion.fragmentId(), _ -> new ArrayList<>()) .addAll(deletion.rowIndexes()); } @@ -1216,7 +1341,8 @@ public void finishMerge( int fragmentId = entry.getKey(); List rowIndexes = entry.getValue(); - log.debug("finishMerge: deleting %d rows from fragment %d, first few indices: %s", + log.debug( + "finishMerge: deleting %d rows from fragment %d, first few indices: %s", rowIndexes.size(), fragmentId, rowIndexes.stream().limit(5).toList()); @@ -1224,8 +1350,10 @@ public void finishMerge( FragmentMetadata updated = dataset.getFragment(fragmentId) .deleteRows(rowIndexes); if (updated != null) { - log.debug("finishMerge: fragment %d updated with deletion vector, deletionFile=%s", - fragmentId, updated.getDeletionFile()); + log.debug( + "finishMerge: fragment %d updated with deletion vector, deletionFile=%s", + fragmentId, + updated.getDeletionFile()); updatedFragments.add(updated); } else { @@ -1237,8 +1365,11 @@ public void finishMerge( Map storageOptions = handle.getStorageOptions(); if (!removedFragmentIds.isEmpty() || !updatedFragments.isEmpty() || !newFragments.isEmpty()) { - log.debug("finishMerge: committing update with %d removed fragments, %d updated fragments, %d new fragments", - removedFragmentIds.size(), updatedFragments.size(), newFragments.size()); + log.debug( + "finishMerge: committing update with %d removed fragments, %d updated fragments, %d new fragments", + removedFragmentIds.size(), + updatedFragments.size(), + newFragments.size()); Update update = Update.builder() .removedFragmentIds(removedFragmentIds) .updatedFragments(updatedFragments) @@ -1266,12 +1397,56 @@ public void finishMerge( throw e; } finally { - dataset.close(); + closeTransactionDataset(transactionId, transactionDataset); } } // ===== Helper Methods ===== + private void registerTransactionDataset(String transactionId, ConnectorSession session, LanceRuntime.DatasetLease datasetLease) + { + TransactionDataset transactionDataset = new TransactionDataset(session.getQueryId(), datasetLease); + TransactionDataset previous = transactionDatasets.putIfAbsent(transactionId, transactionDataset); + if (previous != null) { + closeTransactionDataset(transactionId, transactionDataset); + throw new TrinoException(GENERIC_INTERNAL_ERROR, "Duplicate transaction dataset: " + transactionId); + } + } + + private TransactionDataset removeTransactionDataset(String transactionId) + { + TransactionDataset transactionDataset = transactionDatasets.remove(transactionId); + if (transactionDataset == null) { + throw new TrinoException(GENERIC_INTERNAL_ERROR, "No dataset found for transaction: " + transactionId); + } + return transactionDataset; + } + + private static void closeTransactionDataset(String transactionId, TransactionDataset transactionDataset) + { + try { + transactionDataset.close(); + } + catch (Exception e) { + log.warn(e, "Failed to close dataset for transaction: %s", transactionId); + } + } + + private record TransactionDataset(String queryId, LanceRuntime.DatasetLease datasetLease) + implements AutoCloseable + { + private Dataset dataset() + { + return datasetLease.getDataset(); + } + + @Override + public void close() + { + datasetLease.close(); + } + } + private LanceNamespace getNamespace() { return runtime.getNamespace(); @@ -1486,6 +1661,12 @@ public LanceRuntime getRuntime() return runtime; } + @VisibleForTesting + int getTransactionDatasetCount() + { + return transactionDatasets.size(); + } + /** * Check if the exception is a Lance commit conflict or concurrent modification error. */ @@ -1497,10 +1678,10 @@ private static boolean isCommitConflict(RuntimeException e) String message = current.getMessage(); if (message != null && ( message.toLowerCase().contains("commit conflict") || - message.toLowerCase().contains("concurrent") || - message.toLowerCase().contains("version") || - message.toLowerCase().contains("conflict") || - message.toLowerCase().contains("not found"))) { + message.toLowerCase().contains("concurrent") || + message.toLowerCase().contains("version") || + message.toLowerCase().contains("conflict") || + message.toLowerCase().contains("not found"))) { return true; } // NullPointerException in Fragment/Dataset operations is also a sign of concurrent modification diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageSourceProvider.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageSourceProvider.java index 1d0871b..f3edff9 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageSourceProvider.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageSourceProvider.java @@ -48,8 +48,12 @@ public LancePageSourceProvider(LanceRuntime runtime, LanceConfig lanceConfig) } @Override - public ConnectorPageSource createPageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, - ConnectorSplit split, ConnectorTableHandle tableHandle, List columns, + public ConnectorPageSource createPageSource( + ConnectorTransactionHandle transactionHandle, + ConnectorSession session, + ConnectorSplit split, + ConnectorTableHandle tableHandle, + List columns, DynamicFilter dynamicFilter) { requireNonNull(split, "split is null"); diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageToArrowConverter.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageToArrowConverter.java index 5cfecdb..be7bad1 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageToArrowConverter.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePageToArrowConverter.java @@ -17,6 +17,7 @@ import io.trino.spi.Page; import io.trino.spi.TrinoException; import io.trino.spi.block.Block; +import io.trino.spi.block.RowBlock; import io.trino.spi.connector.ColumnMetadata; import io.trino.spi.type.ArrayType; import io.trino.spi.type.DateType; @@ -56,6 +57,7 @@ import java.util.Map; import java.util.Set; +import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; @@ -239,7 +241,8 @@ public static void validateBlobColumns(List columns, Set if (!(column.getType() instanceof VarbinaryType)) { throw new TrinoException(NOT_SUPPORTED, format("Blob column '%s' must have VARBINARY type, found: %s", - column.getName(), column.getType())); + column.getName(), + column.getType())); } } } @@ -256,13 +259,15 @@ public static void validateVectorColumns(List columns, Map fields = rowType.getFields(); + List fieldBlocks = RowBlock.getRowFieldsFromBlock(block); + if (fieldBlocks.size() != fields.size()) { + throw new TrinoException( + GENERIC_INTERNAL_ERROR, + format("ROW field count mismatch: expected %s, found %s", fields.size(), fieldBlocks.size())); + } + + for (int fieldIndex = 0; fieldIndex < fields.size(); fieldIndex++) { + RowType.Field field = fields.get(fieldIndex); + String fieldName = field.getName().orElse("field" + fieldIndex); + FieldVector childVector = vector.getChild(fieldName); + if (childVector == null) { + throw new TrinoException(GENERIC_INTERNAL_ERROR, format("Missing Arrow child vector for ROW field: %s", fieldName)); + } + writeBlockToVectorAtOffset(fieldBlocks.get(fieldIndex), childVector, field.getType(), rowCount, offset); + } + + for (int i = 0; i < rowCount; i++) { + if (block.isNull(i)) { + vector.setNull(offset + i); + } + else { + vector.setIndexDefined(offset + i); + } + } + vector.setValueCount(offset + rowCount); } } diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePlugin.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePlugin.java index 085312a..012e097 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePlugin.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePlugin.java @@ -34,6 +34,7 @@ public LancePlugin() public LancePlugin(Optional extension) { + LanceRuntime.configureArrowAllocationManager(); this.extension = requireNonNull(extension, "extension is null"); } diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePrefetchingArrowReader.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePrefetchingArrowReader.java index d78dc04..c5382fc 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePrefetchingArrowReader.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LancePrefetchingArrowReader.java @@ -52,35 +52,53 @@ public class LancePrefetchingArrowReader { private static final Logger log = Logger.get(LancePrefetchingArrowReader.class); - /** Default queue depth - number of batches that can be prefetched. */ + /** + * Default queue depth - number of batches that can be prefetched. + */ private static final int DEFAULT_QUEUE_DEPTH = 4; - /** Sentinel batch to signal end of stream. */ + /** + * Sentinel batch to signal end of stream. + */ private static final VectorSchemaRoot END_OF_STREAM = null; private final ArrowReader underlying; private final BufferAllocator allocator; private final int queueDepth; - /** Queue holding prefetched batches ready for consumption. */ + /** + * Queue holding prefetched batches ready for consumption. + */ private final BlockingQueue batchQueue; - /** Background prefetch executor. */ + /** + * Background prefetch executor. + */ private final ExecutorService prefetchExecutor; - /** Current batch being consumed. */ + /** + * Current batch being consumed. + */ private VectorSchemaRoot currentBatch; - /** Flag indicating prefetch thread has finished (either completed or error). */ + /** + * Flag indicating prefetch thread has finished (either completed or error). + */ private final AtomicBoolean prefetchFinished = new AtomicBoolean(false); - /** Error from prefetch thread, if any. */ + /** + * Error from prefetch thread, if any. + */ private final AtomicReference prefetchError = new AtomicReference<>(); - /** Flag indicating consumer has finished. */ + /** + * Flag indicating consumer has finished. + */ private volatile boolean consumerFinished; - /** Total bytes read from Arrow buffers. */ + /** + * Total bytes read from Arrow buffers. + */ private volatile long bytesRead; public LancePrefetchingArrowReader(ArrowReader underlying, BufferAllocator allocator) diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceRuntime.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceRuntime.java index 51f3b6e..d80deaf 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceRuntime.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceRuntime.java @@ -53,6 +53,8 @@ import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -82,6 +84,9 @@ public class LanceRuntime public static final String DEFAULT_SCHEMA = "default"; public static final String TABLE_PATH_SUFFIX = ".lance"; + private static final String ARROW_ALLOCATION_MANAGER_TYPE = "arrow.allocation.manager.type"; + private static final String ARROW_ALLOCATION_MANAGER_TYPE_ENV = "ARROW_ALLOCATION_MANAGER_TYPE"; + private static final String ARROW_ALLOCATION_MANAGER_UNSAFE = "Unsafe"; private static final String ANONYMOUS_USER = "__anonymous__"; // Core resources @@ -95,14 +100,16 @@ public class LanceRuntime private final Map namespaceStorageOptions; // Caches - private final Cache sessionCache; - private final Cache datasetCache; + private final Cache sessionCache; + private final Cache datasetCache; private final Long sessionIndexCacheSizeBytes; private final Long sessionMetadataCacheSizeBytes; @Inject public LanceRuntime(LanceConfig config, @LanceNamespaceProperties Map namespaceProperties) { + configureArrowAllocationManager(); + // Initialize allocator first - it's needed for namespace initialization this.allocator = new RootAllocator( RootAllocator.configBuilder() @@ -153,11 +160,11 @@ public LanceRuntime(LanceConfig config, @LanceNamespaceProperties Map) notification -> { - Session session = notification.getValue(); - if (session != null && !session.isClosed()) { - log.debug("Closing expired session for user: %s", notification.getKey()); - session.close(); + .removalListener((RemovalListener) notification -> { + CachedSession session = notification.getValue(); + if (session != null) { + log.debug("Retiring expired session for user: %s", notification.getKey()); + session.retire(); } }) .build(); @@ -165,21 +172,30 @@ public LanceRuntime(LanceConfig config, @LanceNamespaceProperties Map) notification -> { - Dataset dataset = notification.getValue(); + .removalListener((RemovalListener) notification -> { + CachedDataset dataset = notification.getValue(); if (dataset != null) { - try { - dataset.close(); - } - catch (Exception e) { - log.warn(e, "Failed to close cached dataset"); - } + dataset.retire(); } }) .build(); log.info("LanceRuntime initialized: impl=%s, root=%s, singleLevelNs=%s, maxSessions=%d, maxDatasets=%d", - impl, root, singleLevelNs, config.getCacheSessionMaxEntries(), config.getCacheDatasetMaxEntries()); + impl, + root, + singleLevelNs, + config.getCacheSessionMaxEntries(), + config.getCacheDatasetMaxEntries()); + } + + static void configureArrowAllocationManager() + { + // Arrow 19's Netty allocator is incompatible with Netty 4.2 on Java 25. Prefer + // Unsafe unless the process was explicitly configured with another Arrow allocator. + if (System.getProperty(ARROW_ALLOCATION_MANAGER_TYPE) == null && + System.getenv(ARROW_ALLOCATION_MANAGER_TYPE_ENV) == null) { + System.setProperty(ARROW_ALLOCATION_MANAGER_TYPE, ARROW_ALLOCATION_MANAGER_UNSAFE); + } } // ================== Core Accessors ================== @@ -261,26 +277,36 @@ public List getTableId(String schemaName, String tableName) // ================== Session Management ================== - private Session getOrCreateSession(String userIdentity) + private SessionLease getOrCreateSessionLease(String userIdentity) { String key = normalizeUserIdentity(userIdentity); - try { - return sessionCache.get(key, () -> { - log.debug("Creating new session for user: %s", key); - Session.Builder builder = Session.builder(); - if (sessionIndexCacheSizeBytes != null) { - builder.indexCacheSizeBytes(sessionIndexCacheSizeBytes); - } - if (sessionMetadataCacheSizeBytes != null) { - builder.metadataCacheSizeBytes(sessionMetadataCacheSizeBytes); + while (true) { + try { + CachedSession cachedSession = sessionCache.get(key, () -> new CachedSession(createSession(key))); + SessionLease lease = cachedSession.tryAcquire(); + if (lease != null) { + return lease; } - return builder.build(); - }); + sessionCache.invalidate(key); + } + catch (ExecutionException e) { + log.error(e, "Failed to create session for user: %s", key); + throw new RuntimeException("Failed to create Lance session", e); + } } - catch (ExecutionException e) { - log.error(e, "Failed to create session for user: %s", key); - throw new RuntimeException("Failed to create Lance session", e); + } + + private Session createSession(String key) + { + log.debug("Creating new session for user: %s", key); + Session.Builder builder = Session.builder(); + if (sessionIndexCacheSizeBytes != null) { + builder.indexCacheSizeBytes(sessionIndexCacheSizeBytes); } + if (sessionMetadataCacheSizeBytes != null) { + builder.metadataCacheSizeBytes(sessionMetadataCacheSizeBytes); + } + return builder.build(); } public long getActiveSessionCount() @@ -300,31 +326,76 @@ private static String normalizeUserIdentity(String userIdentity) // ================== Dataset Access ================== - public Dataset getDataset(String userIdentity, String tablePath, Long version, + DatasetLease getDatasetLease( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { - DatasetCacheKey key = new DatasetCacheKey(userIdentity, tablePath, version); + if (!isDatasetCacheable(storageOptions)) { + return openDatasetDirectLease(userIdentity, tablePath, version, storageOptions); + } + + DatasetCacheKey key = new DatasetCacheKey(userIdentity, tablePath, version, storageOptions); + while (true) { + try { + CachedDataset cachedDataset = datasetCache.get( + key, + () -> openCachedDataset(userIdentity, tablePath, version, storageOptions)); + DatasetLease lease = cachedDataset.tryAcquire(); + if (lease != null) { + return lease; + } + datasetCache.invalidate(key); + } + catch (ExecutionException e) { + throw new RuntimeException("Failed to open dataset: " + tablePath, e); + } + } + } + + private CachedDataset openCachedDataset( + String userIdentity, + String tablePath, + Long version, + Map storageOptions) + { + SessionLease sessionLease = getOrCreateSessionLease(userIdentity); try { - return datasetCache.get(key, () -> openDataset(userIdentity, tablePath, version, storageOptions)); + return new CachedDataset(openDataset(sessionLease.getSession(), userIdentity, tablePath, version, storageOptions), sessionLease); } - catch (ExecutionException e) { - throw new RuntimeException("Failed to open dataset: " + tablePath, e); + catch (RuntimeException | Error e) { + sessionLease.close(); + throw e; } } - public Dataset openDatasetDirect(String userIdentity, String tablePath, Long version, + DatasetLease openDatasetDirectLease( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { - return openDataset(userIdentity, tablePath, version, storageOptions); + SessionLease sessionLease = getOrCreateSessionLease(userIdentity); + try { + Dataset dataset = openDataset(sessionLease.getSession(), userIdentity, tablePath, version, storageOptions); + return DatasetLease.direct(dataset, sessionLease); + } + catch (RuntimeException | Error e) { + sessionLease.close(); + throw e; + } } - private Dataset openDataset(String userIdentity, String tablePath, Long version, + private Dataset openDataset( + Session session, + String userIdentity, + String tablePath, + Long version, Map storageOptions) { log.debug("Opening dataset: path=%s, version=%s, user=%s", tablePath, version, userIdentity); - Session session = getOrCreateSession(userIdentity); - ReadOptions.Builder optionsBuilder = new ReadOptions.Builder() .setSession(session); @@ -344,25 +415,31 @@ private Dataset openDataset(String userIdentity, String tablePath, Long version, public long getLatestVersion(String userIdentity, String tablePath, Map storageOptions) { - try (Dataset dataset = openDatasetDirect(userIdentity, tablePath, null, storageOptions)) { - return dataset.version(); + try (DatasetLease datasetLease = openDatasetDirectLease(userIdentity, tablePath, null, storageOptions)) { + return datasetLease.getDataset().version(); } } - public boolean versionExists(String userIdentity, String tablePath, - long version, Map storageOptions) + public boolean versionExists( + String userIdentity, + String tablePath, + long version, + Map storageOptions) { - try (Dataset dataset = openDatasetDirect(userIdentity, tablePath, null, storageOptions)) { - List versions = dataset.listVersions(); + try (DatasetLease datasetLease = openDatasetDirectLease(userIdentity, tablePath, null, storageOptions)) { + List versions = datasetLease.getDataset().listVersions(); return versions.stream().anyMatch(v -> v.getId() == version); } } - public Optional getVersionAtTimestamp(String userIdentity, String tablePath, - long timestampMillis, Map storageOptions) + public Optional getVersionAtTimestamp( + String userIdentity, + String tablePath, + long timestampMillis, + Map storageOptions) { - try (Dataset dataset = openDatasetDirect(userIdentity, tablePath, null, storageOptions)) { - List versions = dataset.listVersions(); + try (DatasetLease datasetLease = openDatasetDirectLease(userIdentity, tablePath, null, storageOptions)) { + List versions = datasetLease.getDataset().listVersions(); Version bestMatch = null; for (Version version : versions) { @@ -375,7 +452,8 @@ public Optional getVersionAtTimestamp(String userIdentity, String tablePat } if (bestMatch != null) { - log.debug("Found version %d at timestamp %s for requested time %s", + log.debug( + "Found version %d at timestamp %s for requested time %s", bestMatch.getId(), bestMatch.getDataTime(), Instant.ofEpochMilli(timestampMillis)); @@ -389,37 +467,57 @@ public Optional getVersionAtTimestamp(String userIdentity, String tablePat // ================== Fragment Access ================== - public List getFragments(String userIdentity, String tablePath, Long version, + public List getFragments( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { - Dataset dataset = getDataset(userIdentity, tablePath, version, storageOptions); - return dataset.getFragments(); + try (DatasetLease datasetLease = getDatasetLease(userIdentity, tablePath, version, storageOptions)) { + return datasetLease.getDataset().getFragments(); + } } - public Fragment getFragment(String userIdentity, String tablePath, Long version, - int fragmentId, Map storageOptions) + public Fragment getFragment( + String userIdentity, + String tablePath, + Long version, + int fragmentId, + Map storageOptions) { - Dataset dataset = getDataset(userIdentity, tablePath, version, storageOptions); - return dataset.getFragment(fragmentId); + try (DatasetLease datasetLease = getDatasetLease(userIdentity, tablePath, version, storageOptions)) { + return datasetLease.getDataset().getFragment(fragmentId); + } } // ================== Schema Access ================== - public Schema getSchema(String userIdentity, String tablePath, Long version, + public Schema getSchema( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { - Dataset dataset = getDataset(userIdentity, tablePath, version, storageOptions); - return dataset.getSchema(); + try (DatasetLease datasetLease = getDatasetLease(userIdentity, tablePath, version, storageOptions)) { + return datasetLease.getDataset().getSchema(); + } } - public LanceSchema getLanceSchema(String userIdentity, String tablePath, Long version, + public LanceSchema getLanceSchema( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { - Dataset dataset = getDataset(userIdentity, tablePath, version, storageOptions); - return dataset.getLanceSchema(); + try (DatasetLease datasetLease = getDatasetLease(userIdentity, tablePath, version, storageOptions)) { + return datasetLease.getDataset().getLanceSchema(); + } } - public Map getColumnHandles(String userIdentity, String tablePath, Long version, + public Map getColumnHandles( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { LanceSchema lanceSchema = getLanceSchema(userIdentity, tablePath, version, storageOptions); @@ -464,7 +562,10 @@ public Map getColumnHandles(String userIdentity, String ta return result; } - public List getColumnHandleList(String userIdentity, String tablePath, Long version, + public List getColumnHandleList( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { LanceSchema lanceSchema = getLanceSchema(userIdentity, tablePath, version, storageOptions); @@ -514,7 +615,10 @@ private static Set getBlobColumnsFromSchema(Schema schema) .collect(Collectors.toSet()); } - public List getColumnMetadata(String userIdentity, String tablePath, Long version, + public List getColumnMetadata( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { Map columnHandles = getColumnHandles(userIdentity, tablePath, version, storageOptions); @@ -523,11 +627,15 @@ public List getColumnMetadata(String userIdentity, String tableP .collect(toImmutableList()); } - public ManifestSummary getManifestSummary(String userIdentity, String tablePath, Long version, + public ManifestSummary getManifestSummary( + String userIdentity, + String tablePath, + Long version, Map storageOptions) { - Dataset dataset = getDataset(userIdentity, tablePath, version, storageOptions); - return dataset.getVersion().getManifestSummary(); + try (DatasetLease datasetLease = getDatasetLease(userIdentity, tablePath, version, storageOptions)) { + return datasetLease.getDataset().getVersion().getManifestSummary(); + } } // ================== Cache Invalidation ================== @@ -538,25 +646,21 @@ public void invalidate(String userIdentity, String tablePath) String normalizedUser = normalizeUserIdentity(userIdentity); - datasetCache.asMap().keySet().removeIf(key -> - Objects.equals(key.userIdentity, normalizedUser) && - Objects.equals(key.tablePath, tablePath)); + List matchingKeys = datasetCache.asMap().keySet().stream() + .filter(key -> Objects.equals(key.userIdentity, normalizedUser) && + Objects.equals(key.tablePath, tablePath)) + .toList(); + datasetCache.invalidateAll(matchingKeys); } // ================== Scanner Operations ================== - public LanceScanner openDatasetScanner(String userIdentity, String tablePath, Long version, - List fragmentIds, ScanOptions scanOptions, Map storageOptions) - { - return openDatasetScanner(userIdentity, tablePath, version, Optional.of(fragmentIds), scanOptions, storageOptions); - } - - public LanceScanner openDatasetScanner(String userIdentity, String tablePath, Long version, - Optional> fragmentIds, ScanOptions scanOptions, Map storageOptions) + LanceScanner openDatasetScanner( + DatasetLease datasetLease, + Optional> fragmentIds, + ScanOptions scanOptions) { - Dataset dataset = getDataset(userIdentity, tablePath, version, storageOptions); - - return dataset.newScan(scanOptionsWithFragmentIds(scanOptions, fragmentIds)); + return datasetLease.getDataset().newScan(scanOptionsWithFragmentIds(scanOptions, fragmentIds)); } private static ScanOptions scanOptionsWithFragmentIds(ScanOptions scanOptions, Optional> fragmentIds) @@ -564,6 +668,8 @@ private static ScanOptions scanOptionsWithFragmentIds(ScanOptions scanOptions, O requireNonNull(scanOptions, "scanOptions is null"); requireNonNull(fragmentIds, "fragmentIds is null"); + // Authoritatively replace the fragment ids on the scan options. An empty value clears any + // fragment filter and scans all fragments, which avoids fragment enumeration for filtered LIMIT queries. return new ScanOptions( fragmentIds.map(List::copyOf), scanOptions.getBatchSize(), @@ -592,26 +698,14 @@ public void close() { log.info("Closing LanceRuntime: %d sessions, %d datasets", sessionCache.size(), datasetCache.size()); - // Close datasets first - datasetCache.asMap().forEach((key, dataset) -> { - if (dataset != null) { - try { - dataset.close(); - } - catch (Exception e) { - log.warn(e, "Failed to close dataset during shutdown"); - } - } - }); + // Retire currently cached datasets before closing sessions they depend on. datasetCache.invalidateAll(); + datasetCache.cleanUp(); // Close sessions - sessionCache.asMap().forEach((key, session) -> { - if (session != null && !session.isClosed()) { - session.close(); - } - }); + sessionCache.asMap().forEach((_, session) -> session.retire()); sessionCache.invalidateAll(); + sessionCache.cleanUp(); // Close namespace if (namespace instanceof Closeable closeable) { @@ -623,9 +717,211 @@ public void close() } } - // Note: We intentionally do NOT close the allocator here. - // Page sources may still be using the allocator asynchronously. - // Arrow allocators just manage memory and will be cleaned up on JVM exit. + try { + allocator.close(); + } + catch (Exception e) { + log.warn(e, "Failed to close Arrow allocator"); + } + } + + private static boolean isDatasetCacheable(Map storageOptions) + { + return storageOptions == null || !storageOptions.containsKey("expires_at_millis"); + } + + private static void closeDataset(Dataset dataset, SessionLease sessionLease) + { + try { + dataset.close(); + } + catch (Exception e) { + log.warn(e, "Failed to close dataset"); + } + finally { + sessionLease.close(); + } + } + + private static final class CachedSession + { + private final Session session; + private final AtomicInteger references = new AtomicInteger(); + private final AtomicBoolean retired = new AtomicBoolean(); + private final AtomicBoolean closed = new AtomicBoolean(); + + private CachedSession(Session session) + { + this.session = session; + } + + private SessionLease tryAcquire() + { + while (true) { + if (retired.get() || closed.get()) { + return null; + } + + int currentReferences = references.get(); + if (references.compareAndSet(currentReferences, currentReferences + 1)) { + if (retired.get() || closed.get()) { + release(); + return null; + } + return new SessionLease(session, this::release); + } + } + } + + private void release() + { + int remainingReferences = references.decrementAndGet(); + if (remainingReferences < 0) { + throw new IllegalStateException("Session reference count is negative"); + } + closeIfUnused(); + } + + private void retire() + { + retired.set(true); + closeIfUnused(); + } + + private void closeIfUnused() + { + if (retired.get() && references.get() == 0 && closed.compareAndSet(false, true) && !session.isClosed()) { + closeSession(session); + } + } + } + + private static final class CachedDataset + { + private final Dataset dataset; + private final SessionLease sessionLease; + private final AtomicInteger references = new AtomicInteger(); + private final AtomicBoolean retired = new AtomicBoolean(); + private final AtomicBoolean closed = new AtomicBoolean(); + + private CachedDataset(Dataset dataset, SessionLease sessionLease) + { + this.dataset = dataset; + this.sessionLease = sessionLease; + } + + private DatasetLease tryAcquire() + { + while (true) { + if (retired.get() || closed.get()) { + return null; + } + + int currentReferences = references.get(); + if (references.compareAndSet(currentReferences, currentReferences + 1)) { + if (retired.get() || closed.get()) { + release(); + return null; + } + return new DatasetLease(dataset, this::release); + } + } + } + + private void release() + { + int remainingReferences = references.decrementAndGet(); + if (remainingReferences < 0) { + throw new IllegalStateException("Dataset reference count is negative"); + } + closeIfUnused(); + } + + private void retire() + { + retired.set(true); + closeIfUnused(); + } + + private void closeIfUnused() + { + if (retired.get() && references.get() == 0 && closed.compareAndSet(false, true)) { + closeDataset(dataset, sessionLease); + } + } + } + + static final class DatasetLease + implements Closeable + { + private final Dataset dataset; + private final Runnable release; + private final AtomicBoolean closed = new AtomicBoolean(); + + private DatasetLease(Dataset dataset, Runnable release) + { + this.dataset = dataset; + this.release = release; + } + + private static DatasetLease direct(Dataset dataset, SessionLease sessionLease) + { + return new DatasetLease(dataset, () -> closeDataset(dataset, sessionLease)); + } + + Dataset getDataset() + { + if (closed.get()) { + throw new IllegalStateException("Dataset lease is closed"); + } + return dataset; + } + + @Override + public void close() + { + if (closed.compareAndSet(false, true)) { + release.run(); + } + } + } + + private static final class SessionLease + implements Closeable + { + private final Session session; + private final Runnable release; + private final AtomicBoolean closed = new AtomicBoolean(); + + private SessionLease(Session session, Runnable release) + { + this.session = session; + this.release = release; + } + + private Session getSession() + { + if (closed.get()) { + throw new IllegalStateException("Session lease is closed"); + } + return session; + } + + @Override + public void close() + { + if (closed.compareAndSet(false, true)) { + release.run(); + } + } + } + + private static Map copyStorageOptions(Map storageOptions) + { + if (storageOptions == null || storageOptions.isEmpty()) { + return Map.of(); + } + return Map.copyOf(storageOptions); } // ================== Cache Key ================== @@ -635,12 +931,14 @@ private static class DatasetCacheKey private final String userIdentity; private final String tablePath; private final Long version; + private final Map storageOptions; - DatasetCacheKey(String userIdentity, String tablePath, Long version) + DatasetCacheKey(String userIdentity, String tablePath, Long version, Map storageOptions) { this.userIdentity = normalizeUserIdentity(userIdentity); this.tablePath = tablePath; this.version = version; + this.storageOptions = copyStorageOptions(storageOptions); } @Override @@ -652,13 +950,26 @@ public boolean equals(Object o) DatasetCacheKey that = (DatasetCacheKey) o; return Objects.equals(userIdentity, that.userIdentity) && Objects.equals(tablePath, that.tablePath) && - Objects.equals(version, that.version); + Objects.equals(version, that.version) && + Objects.equals(storageOptions, that.storageOptions); } @Override public int hashCode() { - return Objects.hash(userIdentity, tablePath, version); + return Objects.hash(userIdentity, tablePath, version, storageOptions); + } + } + + private static void closeSession(Session session) + { + if (!session.isClosed()) { + try { + session.close(); + } + catch (Exception e) { + log.warn(e, "Failed to close session"); + } } } } diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplit.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplit.java index 88fe458..3946657 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplit.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplit.java @@ -17,13 +17,10 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; -import com.google.common.base.Joiner; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import io.trino.spi.connector.ConnectorSplit; import java.util.List; -import java.util.Map; import java.util.Optional; import static com.google.common.base.MoreObjects.toStringHelper; @@ -36,7 +33,6 @@ public class LanceSplit implements ConnectorSplit { - private static final Joiner JOINER = Joiner.on(","); private static final int INSTANCE_SIZE = instanceSize(LanceSplit.class); private final List fragments; @@ -112,18 +108,9 @@ public String toString() .toString(); } - @Override - public Map getSplitInfo() - { - if (allFragments) { - return ImmutableMap.of("fragments", "ALL"); - } - return ImmutableMap.of("fragments", JOINER.join(fragments)); - } - @Override public long getRetainedSizeInBytes() { - return INSTANCE_SIZE + estimatedSizeOf(fragments, e -> sizeOf(Integer.SIZE)); + return INSTANCE_SIZE + estimatedSizeOf(fragments, _ -> sizeOf(Integer.SIZE)); } } diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplitManager.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplitManager.java index ff2470b..68c4dea 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplitManager.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceSplitManager.java @@ -46,8 +46,12 @@ public LanceSplitManager(LanceRuntime runtime) } @Override - public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHandle, ConnectorSession session, - ConnectorTableHandle tableHandle, DynamicFilter dynamicFilter, Constraint constraint) + public ConnectorSplitSource getSplits( + ConnectorTransactionHandle transactionHandle, + ConnectorSession session, + ConnectorTableHandle tableHandle, + DynamicFilter dynamicFilter, + Constraint constraint) { LanceTableHandle lanceTableHandle = (LanceTableHandle) tableHandle; diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableHandle.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableHandle.java index 3915f05..c77cd43 100755 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableHandle.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableHandle.java @@ -81,7 +81,12 @@ public LanceTableHandle( @JsonProperty("countStar") Boolean countStar, @JsonProperty("datasetVersion") Long datasetVersion) { - this(schemaName, tableName, tablePath, tableId, storageOptions, substraitFilter, + this(schemaName, + tableName, + tablePath, + tableId, + storageOptions, + substraitFilter, filterColumns != null ? filterColumns : List.of(), limit != null ? OptionalLong.of(limit) : OptionalLong.empty(), countStar != null && countStar, diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProperties.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProperties.java index cc997c2..379b92c 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProperties.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProperties.java @@ -143,7 +143,7 @@ public static String getFileFormatVersion(Map properties) if (!VALID_STORAGE_VERSIONS.contains(version)) { throw new TrinoException(INVALID_TABLE_PROPERTY, "Invalid file_format_version: '" + value + "'. " + - "Valid values are: legacy, 0.1, 2.0, 2.1, 2.2, stable, next"); + "Valid values are: legacy, 0.1, 2.0, 2.1, 2.2, stable, next"); } return version; } diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/SubstraitExpressionBuilder.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/SubstraitExpressionBuilder.java index 8e267a2..47931dc 100644 --- a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/SubstraitExpressionBuilder.java +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/SubstraitExpressionBuilder.java @@ -117,7 +117,7 @@ public static Optional tupleDomainToSubstrait( .sorted(Comparator.comparingInt(LanceColumnHandle::fieldId)) .toList(); - return Optional.of(serializeAsExtendedExpression(expression.get(), sortedColumns)); + return Optional.of(serializeAsExtendedExpression(expression.orElseThrow(), sortedColumns)); } /** @@ -162,7 +162,7 @@ public static TupleDomainExtractionResult extractTupleDomain( Optional reference = resolveColumnReference(column, columnOrdinals); Optional columnExpr = reference.flatMap(columnReference -> domainToExpression(columnReference, domain)); if (columnExpr.isPresent()) { - columnExpressions.add(columnExpr.get()); + columnExpressions.add(columnExpr.orElseThrow()); pushedDomains.put(column, domain); } else if (!domain.isAll()) { @@ -831,7 +831,7 @@ private static ConnectorExpression extractExpressionsRecursive( // Try to convert entire expression to Substrait Optional substraitExpr = tryConvertToSubstrait(call, assignments, columnOrdinals, columnNames); if (substraitExpr.isPresent()) { - substraitExprs.add(substraitExpr.get()); + substraitExprs.add(substraitExpr.orElseThrow()); return Constant.TRUE; } @@ -914,7 +914,7 @@ private static Optional tryConvertLike( if (patternExpr instanceof Constant constant) { Optional reference = resolveColumnReference(columnExpr, assignments, columnOrdinals); if (reference.isPresent()) { - ColumnReference columnReference = reference.get(); + ColumnReference columnReference = reference.orElseThrow(); Object patternValue = constant.getValue(); if (patternValue instanceof Slice slice && columnReference.column().trinoType() instanceof VarcharType) { String pattern = slice.toStringUtf8(); @@ -943,7 +943,7 @@ private static Optional tryConvertIsNull( Optional reference = resolveColumnReference(args.get(0), assignments, columnOrdinals); if (reference.isPresent()) { - ColumnReference columnReference = reference.get(); + ColumnReference columnReference = reference.orElseThrow(); if (isSupportedType(columnReference.column().trinoType())) { if (!columnNames.contains(columnReference.columnName())) { columnNames.add(columnReference.columnName()); @@ -969,7 +969,7 @@ private static Optional tryConvertNot( if (arg instanceof Call innerCall) { Optional innerExpr = tryConvertToSubstrait(innerCall, assignments, columnOrdinals, columnNames); if (innerExpr.isPresent()) { - return Optional.of(notExpression(innerExpr.get())); + return Optional.of(notExpression(innerExpr.orElseThrow())); } } return Optional.empty(); @@ -998,7 +998,7 @@ private static Optional tryConvertOr( if (converted.isEmpty()) { return Optional.empty(); } - substraitArgs.add(converted.get()); + substraitArgs.add(converted.orElseThrow()); } return Optional.of(orExpressions(substraitArgs)); } @@ -1046,7 +1046,7 @@ private static Optional tryBuildComparison( { Optional reference = resolveColumnReference(columnExpression, assignments, columnOrdinals); if (reference.isPresent()) { - ColumnReference columnReference = reference.get(); + ColumnReference columnReference = reference.orElseThrow(); io.trino.spi.type.Type trinoType = columnReference.column().trinoType(); if (isSupportedType(trinoType)) { Expression fieldRef = fieldReference(columnReference); @@ -1091,7 +1091,7 @@ private static Optional tryConvertIn( return Optional.empty(); } - ColumnReference columnReference = reference.get(); + ColumnReference columnReference = reference.orElseThrow(); LanceColumnHandle lanceColumn = columnReference.column(); io.trino.spi.type.Type trinoType = lanceColumn.trinoType(); if (!isSupportedType(trinoType)) { @@ -1148,7 +1148,7 @@ private static Optional resolveColumnReference( if (target.isEmpty()) { return Optional.empty(); } - ColumnReference targetReference = target.get(); + ColumnReference targetReference = target.orElseThrow(); if (!(fieldDereference.getTarget().getType() instanceof RowType rowType)) { return Optional.empty(); } @@ -1164,7 +1164,7 @@ private static Optional resolveColumnReference( List dereferencePath = new ArrayList<>(targetReference.dereferencePath()); dereferencePath.add(fieldIndex); List dereferenceNames = new ArrayList<>(targetReference.column().dereferenceNames()); - dereferenceNames.add(field.getName().get()); + dereferenceNames.add(field.getName().orElseThrow()); LanceColumnHandle nestedColumn = LanceColumnHandle.nestedColumn( LanceFieldPath.canonicalPath(buildFieldPath(targetReference.column().baseColumnName(), dereferenceNames)), fieldDereference.getType(), diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorSmokeTest.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorSmokeTest.java index 8a894f4..4090d5e 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorSmokeTest.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorSmokeTest.java @@ -56,26 +56,29 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) return switch (connectorBehavior) { // Supported behaviors case SUPPORTS_CREATE_TABLE, - SUPPORTS_CREATE_TABLE_WITH_DATA, - SUPPORTS_INSERT, - SUPPORTS_DELETE, - SUPPORTS_ROW_LEVEL_DELETE, - SUPPORTS_UPDATE, - SUPPORTS_MERGE -> true; + SUPPORTS_CREATE_TABLE_WITH_DATA, + SUPPORTS_INSERT, + SUPPORTS_DELETE, + SUPPORTS_ROW_LEVEL_DELETE, + SUPPORTS_UPDATE, + SUPPORTS_MERGE -> true; // Schema operations - depends on namespace configuration case SUPPORTS_CREATE_SCHEMA -> getNamespaceTestConfig().supportsCreateSchema(); // Not supported behaviors case SUPPORTS_DROP_SCHEMA_CASCADE, - SUPPORTS_RENAME_SCHEMA, - SUPPORTS_RENAME_TABLE, - SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, - SUPPORTS_TRUNCATE, - SUPPORTS_CREATE_VIEW, - SUPPORTS_COMMENT_ON_VIEW_COLUMN, - SUPPORTS_CREATE_MATERIALIZED_VIEW, - SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; + SUPPORTS_RENAME_SCHEMA, + SUPPORTS_RENAME_TABLE, + SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, + SUPPORTS_DEFAULT_COLUMN_VALUE, + SUPPORTS_SET_DEFAULT_COLUMN_VALUE, + SUPPORTS_DROP_DEFAULT_COLUMN_VALUE, + SUPPORTS_TRUNCATE, + SUPPORTS_CREATE_VIEW, + SUPPORTS_COMMENT_ON_VIEW_COLUMN, + SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; default -> super.hasBehavior(connectorBehavior); }; diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorTest.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorTest.java index 6f112ca..1b4da4c 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorTest.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/BaseLanceConnectorTest.java @@ -72,53 +72,58 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) return switch (connectorBehavior) { // Supported write behaviors case SUPPORTS_CREATE_TABLE, - SUPPORTS_CREATE_TABLE_WITH_DATA, - SUPPORTS_CREATE_OR_REPLACE_TABLE, - SUPPORTS_INSERT -> true; + SUPPORTS_CREATE_TABLE_WITH_DATA, + SUPPORTS_CREATE_OR_REPLACE_TABLE, + SUPPORTS_INSERT -> true; // Complex types - ROW and MAP not fully supported for writes case SUPPORTS_ROW_TYPE, - SUPPORTS_MAP_TYPE -> false; + SUPPORTS_MAP_TYPE -> false; // Schema operations - depends on namespace configuration case SUPPORTS_CREATE_SCHEMA -> getNamespaceTestConfig().supportsCreateSchema(); case SUPPORTS_DROP_SCHEMA_CASCADE, - SUPPORTS_RENAME_SCHEMA -> false; + SUPPORTS_RENAME_SCHEMA -> false; // Table modification operations - not supported case SUPPORTS_RENAME_TABLE, - SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, - SUPPORTS_ADD_COLUMN, - SUPPORTS_ADD_COLUMN_WITH_COMMENT, - SUPPORTS_ADD_COLUMN_NOT_NULL_CONSTRAINT, - SUPPORTS_DROP_COLUMN, - SUPPORTS_RENAME_COLUMN, - SUPPORTS_SET_COLUMN_TYPE -> false; + SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, + SUPPORTS_ADD_COLUMN, + SUPPORTS_ADD_COLUMN_WITH_COMMENT, + SUPPORTS_ADD_COLUMN_NOT_NULL_CONSTRAINT, + SUPPORTS_DROP_COLUMN, + SUPPORTS_RENAME_COLUMN, + SUPPORTS_SET_COLUMN_TYPE -> false; + + case SUPPORTS_DEFAULT_COLUMN_VALUE, + SUPPORTS_SET_DEFAULT_COLUMN_VALUE, + SUPPORTS_DROP_DEFAULT_COLUMN_VALUE -> false; // Row-level modification operations - supported via merge-on-read case SUPPORTS_DELETE, - SUPPORTS_ROW_LEVEL_DELETE, - SUPPORTS_UPDATE, - SUPPORTS_MERGE -> true; + SUPPORTS_ROW_LEVEL_DELETE, + SUPPORTS_UPDATE, + SUPPORTS_MERGE -> true; // Truncate not yet supported case SUPPORTS_TRUNCATE -> false; // View operations - not supported case SUPPORTS_CREATE_VIEW, - SUPPORTS_COMMENT_ON_VIEW_COLUMN, - SUPPORTS_CREATE_MATERIALIZED_VIEW, - SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; + SUPPORTS_COMMENT_ON_VIEW_COLUMN, + SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; // Comment operations - not supported case SUPPORTS_COMMENT_ON_TABLE, - SUPPORTS_COMMENT_ON_COLUMN -> false; + SUPPORTS_COMMENT_ON_COLUMN -> false; // Constraint operations - not supported case SUPPORTS_NOT_NULL_CONSTRAINT -> false; // Pushdown operations - not currently supported - case SUPPORTS_TOPN_PUSHDOWN -> false; + case SUPPORTS_LIMIT_PUSHDOWN, + SUPPORTS_TOPN_PUSHDOWN -> false; // Date handling - negative dates may not be supported case SUPPORTS_NEGATIVE_DATE -> false; @@ -468,11 +473,11 @@ protected void verifyVersionedQueryFailurePermissible(Exception e) // Lance supports time travel, so we expect specific error messages instead of "not supported" assertThat(e).hasMessageMatching( "Lance connector does not support start version for time travel|" + - "Lance version number must be positive: .*|" + - "Lance version does not exist: .*|" + - "Unsupported type for Lance version: .*\\..*|" + - "Unsupported type for Lance temporal version: .*|" + - "No Lance version found at or before timestamp: .*"); + "Lance version number must be positive: .*|" + + "Lance version does not exist: .*|" + + "Unsupported type for Lance version: .*\\..*|" + + "Unsupported type for Lance temporal version: .*|" + + "No Lance version found at or before timestamp: .*"); } // ===== Time Travel Tests ===== diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/LanceQueryRunner.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/LanceQueryRunner.java index 6b2e36a..0014cd6 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/LanceQueryRunner.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/LanceQueryRunner.java @@ -256,7 +256,7 @@ public DistributedQueryRunner build() { // Apply namespace test config if specified if (namespaceTestConfig.isPresent()) { - LanceNamespaceTestConfig config = namespaceTestConfig.get(); + LanceNamespaceTestConfig config = namespaceTestConfig.orElseThrow(); // Only apply if not S3 config (S3 properties are set separately in builderForS3) if (!config.isS3Config() && useTempDirectory) { Path tempDir = Files.createTempDirectory("lance-trino-test"); @@ -277,7 +277,7 @@ else if (useTempDirectory && !connectorProperties.containsKey("lance.uri")) { } // Create S3 bucket if using S3 configuration - if (namespaceTestConfig.isPresent() && namespaceTestConfig.get().isS3Config()) { + if (namespaceTestConfig.isPresent() && namespaceTestConfig.orElseThrow().isS3Config()) { String endpoint = connectorProperties.get("lance.storage.aws_endpoint"); String root = connectorProperties.get("lance.root"); if (endpoint != null && root != null && root.startsWith("s3://")) { @@ -299,14 +299,14 @@ else if (useTempDirectory && !connectorProperties.containsKey("lance.uri")) { Boolean.parseBoolean(connectorProperties.get("lance.single_level_ns")); // Check if we have a parent namespace configuration - boolean hasParent = namespaceTestConfig.isPresent() && namespaceTestConfig.get().hasParent(); + boolean hasParent = namespaceTestConfig.isPresent() && namespaceTestConfig.orElseThrow().hasParent(); if (hasParent) { // For configurations with parent namespaces, we need to create the parent // namespaces first. We do this by temporarily creating the catalog without // the parent setting, creating the parent namespaces, then recreating with // the parent setting. - LanceNamespaceTestConfig config = namespaceTestConfig.get(); + LanceNamespaceTestConfig config = namespaceTestConfig.orElseThrow(); List parentLevels = config.getParentLevels(); log.info("Creating parent namespaces: %s", parentLevels); diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestArrowAllocationManager.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestArrowAllocationManager.java new file mode 100644 index 0000000..09eb3cd --- /dev/null +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestArrowAllocationManager.java @@ -0,0 +1,81 @@ +/* + * Licensed 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 io.trino.plugin.lance; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.parallel.ResourceLock; +import org.junit.jupiter.api.parallel.Resources; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assumptions.assumeThat; + +public class TestArrowAllocationManager +{ + private static final String ARROW_ALLOCATION_MANAGER_TYPE = "arrow.allocation.manager.type"; + + @Test + public void testDirectRootAllocator() + { + try (BufferAllocator allocator = new RootAllocator(); + ArrowBuf buffer = allocator.buffer(Long.BYTES)) { + buffer.setLong(0, 42); + assertThat(buffer.getLong(0)).isEqualTo(42); + } + } + + // Exercises LanceRuntime.configureArrowAllocationManager() directly. It mutates the global + // arrow.allocation.manager.type system property, so it locks SYSTEM_PROPERTIES to avoid racing + // other property-sensitive tests. + @Test + @ResourceLock(Resources.SYSTEM_PROPERTIES) + public void testConfigureArrowAllocationManager() + { + // The method is a no-op when the env var is set; skip if an operator override is present. + assumeThat(System.getenv("ARROW_ALLOCATION_MANAGER_TYPE")) + .as("operator override via environment variable") + .isNull(); + + // Force Arrow's allocation-manager factory to resolve now (it is read once, lazily), so + // clearing the property below cannot affect a concurrently-created allocator. + try (BufferAllocator allocator = new RootAllocator(); + ArrowBuf buffer = allocator.buffer(Long.BYTES)) { + assertThat(buffer.capacity()).isGreaterThanOrEqualTo(Long.BYTES); + } + + String original = System.getProperty(ARROW_ALLOCATION_MANAGER_TYPE); + try { + // When unset, the connector pins the Arrow "Unsafe" allocator because Arrow 19's Netty + // allocator is incompatible with Netty 4.2 on Java 25. + System.clearProperty(ARROW_ALLOCATION_MANAGER_TYPE); + LanceRuntime.configureArrowAllocationManager(); + assertThat(System.getProperty(ARROW_ALLOCATION_MANAGER_TYPE)).isEqualTo("Unsafe"); + + // An explicit override must be left untouched. + System.setProperty(ARROW_ALLOCATION_MANAGER_TYPE, "Netty"); + LanceRuntime.configureArrowAllocationManager(); + assertThat(System.getProperty(ARROW_ALLOCATION_MANAGER_TYPE)).isEqualTo("Netty"); + } + finally { + if (original == null) { + System.clearProperty(ARROW_ALLOCATION_MANAGER_TYPE); + } + else { + System.setProperty(ARROW_ALLOCATION_MANAGER_TYPE, original); + } + } + } +} diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceArrowToPageScanner.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceArrowToPageScanner.java index 16306af..ac1aefd 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceArrowToPageScanner.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceArrowToPageScanner.java @@ -31,6 +31,7 @@ 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.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; @@ -122,13 +123,21 @@ public void setUp() try (LanceFragmentPageSource pageSource = new LanceFragmentPageSource( lanceTableHandle, columns, lanceSplit.getFragments(), Collections.emptyMap(), 8192, null, runtime)) { - page = pageSource.getNextPage(); + page = pageSource.getNextSourcePage().getPage(); } assertThat(page).isNotNull(); assertThat(page.getPositionCount()).isEqualTo(2); } + @AfterEach + public void tearDown() + { + if (runtime != null) { + runtime.close(); + } + } + @Test public void testBigint() { diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceBlobEncoding.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceBlobEncoding.java index 25c7c07..a2d6f98 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceBlobEncoding.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceBlobEncoding.java @@ -57,7 +57,8 @@ public void testCreateTableAsSelectWithBlobEncoding() String tableName = "test_blob_ctas_" + System.currentTimeMillis(); try { // Create table with blob encoding via CTAS - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (blob_columns = 'content') AS " + "SELECT CAST(1 AS BIGINT) as id, X'48454C4C4F' as content", 1); @@ -77,7 +78,8 @@ public void testInsertIntoBlobColumn() String tableName = "test_blob_insert_" + System.currentTimeMillis(); try { // Create table with blob encoding - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (blob_columns = 'content') AS " + "SELECT CAST(1 AS BIGINT) as id, X'48454C4C4F' as content", 1); @@ -110,7 +112,8 @@ public void testMultipleBlobColumns() String tableName = "test_multi_blob_" + System.currentTimeMillis(); try { // Create table with multiple blob columns - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (blob_columns = 'content1, content2') AS " + "SELECT CAST(1 AS BIGINT) as id, X'41' as content1, X'42' as content2", 1); @@ -129,7 +132,8 @@ public void testBlobWithMixedColumns() String tableName = "test_blob_mixed_" + System.currentTimeMillis(); try { // Create table with mix of blob and non-blob columns - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (blob_columns = 'blob_content') AS " + "SELECT CAST(1 AS BIGINT) as id, " + " 'regular text' as text_content, " + @@ -152,7 +156,8 @@ public void testBlobVirtualColumnsSelectable() String tableName = "test_blob_virtual_select_" + System.currentTimeMillis(); try { // Create table with blob encoding - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (blob_columns = 'content') AS " + "SELECT CAST(1 AS BIGINT) as id, X'48454C4C4F' as content", 1); @@ -178,7 +183,8 @@ public void testMultipleBlobVirtualColumns() String tableName = "test_multi_blob_virtual_" + System.currentTimeMillis(); try { // Create table with multiple blob columns - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (blob_columns = 'data1, data2') AS " + "SELECT CAST(1 AS BIGINT) as id, X'41' as data1, X'424344' as data2", 1); diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorSmokeTest.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorSmokeTest.java index 0da97d7..8a43a93 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorSmokeTest.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorSmokeTest.java @@ -61,25 +61,25 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) return switch (connectorBehavior) { // Supported behaviors case SUPPORTS_CREATE_TABLE, - SUPPORTS_CREATE_TABLE_WITH_DATA, - SUPPORTS_INSERT, - SUPPORTS_DELETE, - SUPPORTS_ROW_LEVEL_DELETE, - SUPPORTS_UPDATE, - SUPPORTS_MERGE -> true; + SUPPORTS_CREATE_TABLE_WITH_DATA, + SUPPORTS_INSERT, + SUPPORTS_DELETE, + SUPPORTS_ROW_LEVEL_DELETE, + SUPPORTS_UPDATE, + SUPPORTS_MERGE -> true; // Not supported behaviors // CASCADE is not supported for DROP SCHEMA case SUPPORTS_CREATE_SCHEMA, - SUPPORTS_DROP_SCHEMA_CASCADE, - SUPPORTS_RENAME_SCHEMA, - SUPPORTS_RENAME_TABLE, - SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, - SUPPORTS_TRUNCATE, - SUPPORTS_CREATE_VIEW, - SUPPORTS_COMMENT_ON_VIEW_COLUMN, - SUPPORTS_CREATE_MATERIALIZED_VIEW, - SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; + SUPPORTS_DROP_SCHEMA_CASCADE, + SUPPORTS_RENAME_SCHEMA, + SUPPORTS_RENAME_TABLE, + SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, + SUPPORTS_TRUNCATE, + SUPPORTS_CREATE_VIEW, + SUPPORTS_COMMENT_ON_VIEW_COLUMN, + SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; default -> super.hasBehavior(connectorBehavior); }; diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorTest.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorTest.java index ca19f83..ee33a59 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorTest.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceConnectorTest.java @@ -117,52 +117,57 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) return switch (connectorBehavior) { // Supported write behaviors case SUPPORTS_CREATE_TABLE, - SUPPORTS_CREATE_TABLE_WITH_DATA, - SUPPORTS_CREATE_OR_REPLACE_TABLE, - SUPPORTS_INSERT, - SUPPORTS_DELETE, - SUPPORTS_ROW_LEVEL_DELETE, - SUPPORTS_UPDATE, - SUPPORTS_MERGE -> true; + SUPPORTS_CREATE_TABLE_WITH_DATA, + SUPPORTS_CREATE_OR_REPLACE_TABLE, + SUPPORTS_INSERT, + SUPPORTS_DELETE, + SUPPORTS_ROW_LEVEL_DELETE, + SUPPORTS_UPDATE, + SUPPORTS_MERGE -> true; // Complex types - ROW and MAP not fully supported for writes case SUPPORTS_ROW_TYPE, - SUPPORTS_MAP_TYPE -> false; + SUPPORTS_MAP_TYPE -> false; // Schema operations - not supported in single-level mode (which builderForWriteTests uses) // CASCADE is not supported for DROP SCHEMA case SUPPORTS_CREATE_SCHEMA, - SUPPORTS_RENAME_SCHEMA, - SUPPORTS_DROP_SCHEMA_CASCADE -> false; + SUPPORTS_RENAME_SCHEMA, + SUPPORTS_DROP_SCHEMA_CASCADE -> false; // Table modification operations - not supported case SUPPORTS_RENAME_TABLE, - SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, - SUPPORTS_ADD_COLUMN, - SUPPORTS_ADD_COLUMN_WITH_COMMENT, - SUPPORTS_ADD_COLUMN_NOT_NULL_CONSTRAINT, - SUPPORTS_DROP_COLUMN, - SUPPORTS_RENAME_COLUMN, - SUPPORTS_SET_COLUMN_TYPE -> false; + SUPPORTS_RENAME_TABLE_ACROSS_SCHEMAS, + SUPPORTS_ADD_COLUMN, + SUPPORTS_ADD_COLUMN_WITH_COMMENT, + SUPPORTS_ADD_COLUMN_NOT_NULL_CONSTRAINT, + SUPPORTS_DROP_COLUMN, + SUPPORTS_RENAME_COLUMN, + SUPPORTS_SET_COLUMN_TYPE -> false; + + case SUPPORTS_DEFAULT_COLUMN_VALUE, + SUPPORTS_SET_DEFAULT_COLUMN_VALUE, + SUPPORTS_DROP_DEFAULT_COLUMN_VALUE -> false; // Truncate is not supported case SUPPORTS_TRUNCATE -> false; // View operations - not supported case SUPPORTS_CREATE_VIEW, - SUPPORTS_COMMENT_ON_VIEW_COLUMN, - SUPPORTS_CREATE_MATERIALIZED_VIEW, - SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; + SUPPORTS_COMMENT_ON_VIEW_COLUMN, + SUPPORTS_CREATE_MATERIALIZED_VIEW, + SUPPORTS_COMMENT_ON_MATERIALIZED_VIEW_COLUMN -> false; // Comment operations - not supported case SUPPORTS_COMMENT_ON_TABLE, - SUPPORTS_COMMENT_ON_COLUMN -> false; + SUPPORTS_COMMENT_ON_COLUMN -> false; // Constraint operations - not supported case SUPPORTS_NOT_NULL_CONSTRAINT -> false; // Pushdown operations - not currently supported - case SUPPORTS_TOPN_PUSHDOWN -> false; + case SUPPORTS_LIMIT_PUSHDOWN, + SUPPORTS_TOPN_PUSHDOWN -> false; // Date handling - negative dates may not be supported case SUPPORTS_NEGATIVE_DATE -> false; @@ -184,11 +189,11 @@ protected void verifyVersionedQueryFailurePermissible(Exception e) // Lance supports time travel, so we expect specific error messages instead of "not supported" assertThat(e).hasMessageMatching( "Lance connector does not support start version for time travel|" + - "Lance version number must be positive: .*|" + - "Lance version does not exist: .*|" + - "Unsupported type for Lance version: .*\\..*|" + - "Unsupported type for Lance temporal version: .*|" + - "No Lance version found at or before timestamp: .*"); + "Lance version number must be positive: .*|" + + "Lance version does not exist: .*|" + + "Unsupported type for Lance version: .*\\..*|" + + "Unsupported type for Lance temporal version: .*|" + + "No Lance version found at or before timestamp: .*"); } @Test @@ -511,7 +516,10 @@ public void testReadLargeUtf8Dataset(@TempDir Path tempDir) // Step 1: Create a Lance dataset with LargeUtf8 column using the Java SDK try (BufferAllocator allocator = new RootAllocator()) { // Create empty dataset with schema - Dataset dataset = Dataset.create(allocator, datasetPath, LARGE_UTF8_SCHEMA, + Dataset dataset = Dataset.create( + allocator, + datasetPath, + LARGE_UTF8_SCHEMA, new WriteParams.Builder().build()); dataset.close(); @@ -545,33 +553,33 @@ public void testReadLargeUtf8Dataset(@TempDir Path tempDir) // Use single-level mode since we're using a flat directory structure LanceConfig config = new LanceConfig().setSingleLevelNs(true); Map catalogProperties = ImmutableMap.of("lance.root", tempDir.toString()); - LanceRuntime runtime = new LanceRuntime(config, catalogProperties); - - // Get column handles using the table path - String tablePath = datasetPath; - Map columnHandles = runtime.getColumnHandles(null, tablePath, null, Map.of()); - assertThat(columnHandles).hasSize(2); - - // Verify the large_text column is mapped to VARCHAR - LanceColumnHandle largeTextHandle = (LanceColumnHandle) columnHandles.get("large_text"); - assertThat(largeTextHandle).isNotNull(); - assertThat(largeTextHandle.trinoType()).isEqualTo(VarcharType.VARCHAR); - - // Verify id column - LanceColumnHandle idHandle = (LanceColumnHandle) columnHandles.get("id"); - assertThat(idHandle).isNotNull(); - assertThat(idHandle.trinoType()).isEqualTo(INTEGER); - - // Step 3: Verify table metadata using the table path - List columnsMetadata = runtime.getColumnMetadata(null, tablePath, null, Map.of()); - assertThat(columnsMetadata).hasSize(2); - - // Find the large_text column metadata - ColumnMetadata largeTextMetadata = columnsMetadata.stream() - .filter(cm -> cm.getName().equals("large_text")) - .findFirst() - .orElseThrow(); - assertThat(largeTextMetadata.getType()).isEqualTo(VarcharType.VARCHAR); + try (LanceRuntime runtime = new LanceRuntime(config, catalogProperties)) { + // Get column handles using the table path + String tablePath = datasetPath; + Map columnHandles = runtime.getColumnHandles(null, tablePath, null, Map.of()); + assertThat(columnHandles).hasSize(2); + + // Verify the large_text column is mapped to VARCHAR + LanceColumnHandle largeTextHandle = (LanceColumnHandle) columnHandles.get("large_text"); + assertThat(largeTextHandle).isNotNull(); + assertThat(largeTextHandle.trinoType()).isEqualTo(VarcharType.VARCHAR); + + // Verify id column + LanceColumnHandle idHandle = (LanceColumnHandle) columnHandles.get("id"); + assertThat(idHandle).isNotNull(); + assertThat(idHandle.trinoType()).isEqualTo(INTEGER); + + // Step 3: Verify table metadata using the table path + List columnsMetadata = runtime.getColumnMetadata(null, tablePath, null, Map.of()); + assertThat(columnsMetadata).hasSize(2); + + // Find the large_text column metadata + ColumnMetadata largeTextMetadata = columnsMetadata.stream() + .filter(cm -> cm.getName().equals("large_text")) + .findFirst() + .orElseThrow(); + assertThat(largeTextMetadata.getType()).isEqualTo(VarcharType.VARCHAR); + } } } @@ -582,7 +590,10 @@ public void testLanceMetadataWithLargeUtf8(@TempDir Path tempDir) try (BufferAllocator allocator = new RootAllocator()) { // Create and populate dataset - Dataset dataset = Dataset.create(allocator, datasetPath, LARGE_UTF8_SCHEMA, + Dataset dataset = Dataset.create( + allocator, + datasetPath, + LARGE_UTF8_SCHEMA, new WriteParams.Builder().build()); dataset.close(); @@ -595,7 +606,10 @@ public void testLanceMetadataWithLargeUtf8(@TempDir Path tempDir) textVector.setSafe(0, "test".getBytes(StandardCharsets.UTF_8)); root.setRowCount(1); - List fragments = Fragment.create(datasetPath, allocator, root, + List fragments = Fragment.create( + datasetPath, + allocator, + root, new WriteParams.Builder().build()); FragmentOperation.Append appendOp = new FragmentOperation.Append(fragments); Dataset.commit(allocator, datasetPath, appendOp, Optional.of(1L)).close(); @@ -605,25 +619,26 @@ public void testLanceMetadataWithLargeUtf8(@TempDir Path tempDir) // Use single-level mode since we're using a flat directory structure LanceConfig config = new LanceConfig().setSingleLevelNs(true); Map catalogProperties = ImmutableMap.of("lance.root", tempDir.toString()); - LanceRuntime runtime = new LanceRuntime(config, catalogProperties); - JsonCodec commitTaskDataCodec = JsonCodec.jsonCodec(LanceCommitTaskData.class); - JsonCodec mergeCommitDataCodec = JsonCodec.jsonCodec(LanceMergeCommitData.class); - LanceMetadata metadata = new LanceMetadata(runtime, config, commitTaskDataCodec, mergeCommitDataCodec); - - // Get table handle - this should NOT return null anymore - LanceTableHandle tableHandle = (LanceTableHandle) metadata.getTableHandle( - SESSION, - new SchemaTableName("default", "metadata_test"), - Optional.empty(), - Optional.empty()); - assertThat(tableHandle).isNotNull(); - - // Get table metadata - this should NOT return null anymore - var tableMetadata = metadata.getTableMetadata(SESSION, tableHandle); - assertThat(tableMetadata) - .describedAs("getTableMetadata should not return null for LargeUtf8 columns") - .isNotNull(); - assertThat(tableMetadata.getColumns()).hasSize(2); + try (LanceRuntime runtime = new LanceRuntime(config, catalogProperties)) { + JsonCodec commitTaskDataCodec = JsonCodec.jsonCodec(LanceCommitTaskData.class); + JsonCodec mergeCommitDataCodec = JsonCodec.jsonCodec(LanceMergeCommitData.class); + LanceMetadata metadata = new LanceMetadata(runtime, config, commitTaskDataCodec, mergeCommitDataCodec); + + // Get table handle - this should NOT return null anymore + LanceTableHandle tableHandle = (LanceTableHandle) metadata.getTableHandle( + SESSION, + new SchemaTableName("default", "metadata_test"), + Optional.empty(), + Optional.empty()); + assertThat(tableHandle).isNotNull(); + + // Get table metadata - this should NOT return null anymore + var tableMetadata = metadata.getTableMetadata(SESSION, tableHandle); + assertThat(tableMetadata) + .describedAs("getTableMetadata should not return null for LargeUtf8 columns") + .isNotNull(); + assertThat(tableMetadata.getColumns()).hasSize(2); + } } } @@ -647,7 +662,10 @@ public void testReadUInt4Dataset(@TempDir Path tempDir) try (BufferAllocator allocator = new RootAllocator()) { // Create empty dataset with schema - Dataset dataset = Dataset.create(allocator, datasetPath, UINT32_SCHEMA, + Dataset dataset = Dataset.create( + allocator, + datasetPath, + UINT32_SCHEMA, new WriteParams.Builder().build()); dataset.close(); @@ -681,45 +699,49 @@ public void testReadUInt4Dataset(@TempDir Path tempDir) // Read back and verify schema mapping LanceConfig config = new LanceConfig().setSingleLevelNs(true); Map catalogProperties = ImmutableMap.of("lance.root", tempDir.toString()); - LanceRuntime runtime = new LanceRuntime(config, catalogProperties); - - // Verify the unsigned_val column is mapped to BIGINT - Map columnHandles = runtime.getColumnHandles(null, datasetPath, null, Map.of()); - assertThat(columnHandles).hasSize(2); - - LanceColumnHandle unsignedHandle = (LanceColumnHandle) columnHandles.get("unsigned_val"); - assertThat(unsignedHandle).isNotNull(); - assertThat(unsignedHandle.trinoType()).isEqualTo(BIGINT); - - // Read data through the page source and verify values - LanceTableHandle tableHandle = new LanceTableHandle("default", "uint4_test", - datasetPath, List.of("uint4_test"), Map.of()); - LanceSplitManager splitManager = new LanceSplitManager(runtime); - var splitSource = splitManager.getSplits(null, SESSION, tableHandle, null, null); - var batch = splitSource.getNextBatch(10).get(); - LanceSplit split = (LanceSplit) batch.getSplits().get(0); - - List columns = runtime.getColumnHandleList(null, datasetPath, null, Map.of()); - try (LanceFragmentPageSource pageSource = new LanceFragmentPageSource( - tableHandle, columns, split.getFragments(), Map.of(), 8192, null, runtime)) { - io.trino.spi.Page page = pageSource.getNextPage(); - assertThat(page).isNotNull(); - assertThat(page.getPositionCount()).isEqualTo(2); - - // Find the unsigned_val column index - int unsignedIdx = -1; - for (int i = 0; i < columns.size(); i++) { - if (columns.get(i).name().equals("unsigned_val")) { - unsignedIdx = i; - break; + try (LanceRuntime runtime = new LanceRuntime(config, catalogProperties)) { + // Verify the unsigned_val column is mapped to BIGINT + Map columnHandles = runtime.getColumnHandles(null, datasetPath, null, Map.of()); + assertThat(columnHandles).hasSize(2); + + LanceColumnHandle unsignedHandle = (LanceColumnHandle) columnHandles.get("unsigned_val"); + assertThat(unsignedHandle).isNotNull(); + assertThat(unsignedHandle.trinoType()).isEqualTo(BIGINT); + + // Read data through the page source and verify values + LanceTableHandle tableHandle = new LanceTableHandle( + "default", + "uint4_test", + datasetPath, + List.of("uint4_test"), + Map.of()); + LanceSplitManager splitManager = new LanceSplitManager(runtime); + var splitSource = splitManager.getSplits(null, SESSION, tableHandle, null, null); + var batch = splitSource.getNextBatch(10).get(); + LanceSplit split = (LanceSplit) batch.getSplits().get(0); + + List columns = runtime.getColumnHandleList(null, datasetPath, null, Map.of()); + try (LanceFragmentPageSource pageSource = new LanceFragmentPageSource( + tableHandle, columns, split.getFragments(), Map.of(), 8192, null, runtime)) { + io.trino.spi.Page page = pageSource.getNextSourcePage().getPage(); + assertThat(page).isNotNull(); + assertThat(page.getPositionCount()).isEqualTo(2); + + // Find the unsigned_val column index + int unsignedIdx = -1; + for (int i = 0; i < columns.size(); i++) { + if (columns.get(i).name().equals("unsigned_val")) { + unsignedIdx = i; + break; + } } - } - assertThat(unsignedIdx).isGreaterThanOrEqualTo(0); + assertThat(unsignedIdx).isGreaterThanOrEqualTo(0); - // Small value should read correctly - assertThat(BIGINT.getLong(page.getBlock(unsignedIdx), 0)).isEqualTo(42L); - // Value exceeding Integer.MAX_VALUE should be correctly promoted to unsigned long - assertThat(BIGINT.getLong(page.getBlock(unsignedIdx), 1)).isEqualTo(3_000_000_000L); + // Small value should read correctly + assertThat(BIGINT.getLong(page.getBlock(unsignedIdx), 0)).isEqualTo(42L); + // Value exceeding Integer.MAX_VALUE should be correctly promoted to unsigned long + assertThat(BIGINT.getLong(page.getBlock(unsignedIdx), 1)).isEqualTo(3_000_000_000L); + } } } } diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceCountPageSource.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceCountPageSource.java index 5640fd3..91282ac 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceCountPageSource.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceCountPageSource.java @@ -20,6 +20,7 @@ import io.trino.spi.connector.ConnectorTableHandle; import io.trino.spi.connector.SchemaTableName; import io.trino.testing.TestingConnectorSession; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; @@ -58,6 +59,14 @@ public void setUp() this.metadata = new LanceMetadata(runtime, lanceConfig, commitTaskDataCodec, mergeCommitDataCodec); } + @AfterEach + public void tearDown() + { + if (runtime != null) { + runtime.close(); + } + } + @Test public void testCountStarWithoutFilter() { @@ -75,7 +84,7 @@ public void testCountStarWithoutFilter() runtime)) { assertThat(pageSource.isFinished()).isFalse(); - Page page = pageSource.getNextPage(); + Page page = pageSource.getNextSourcePage().getPage(); assertThat(page).isNotNull(); assertThat(page.getChannelCount()).isEqualTo(1); assertThat(page.getPositionCount()).isEqualTo(1); @@ -85,7 +94,7 @@ public void testCountStarWithoutFilter() assertThat(count).isEqualTo(4L); // Second call should return null - assertThat(pageSource.getNextPage()).isNull(); + assertThat(pageSource.getNextSourcePage()).isNull(); assertThat(pageSource.isFinished()).isTrue(); } } diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceFragmentPageSource.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceFragmentPageSource.java index 2e367e6..d18eb07 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceFragmentPageSource.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceFragmentPageSource.java @@ -28,6 +28,7 @@ import io.trino.spi.predicate.ValueSet; import io.trino.testing.TestingConnectorSession; import org.apache.arrow.vector.ipc.ArrowReader; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; @@ -76,6 +77,14 @@ public void setUp() this.splitManager = new LanceSplitManager(runtime); } + @AfterEach + public void tearDown() + { + if (runtime != null) { + runtime.close(); + } + } + @Test public void testDatasetScanWithoutFragmentIdsRespectsLimit() throws Exception @@ -86,13 +95,12 @@ public void testDatasetScanWithoutFragmentIdsRespectsLimit() .limit(3) .build(); - try (LanceScanner scanner = runtime.openDatasetScanner( + try (LanceRuntime.DatasetLease datasetLease = runtime.getDatasetLease( TestingConnectorSession.SESSION.getUser(), tableHandle.getTablePath(), tableHandle.getDatasetVersion(), - Optional.empty(), - scanOptions, - Collections.emptyMap())) { + Collections.emptyMap()); + LanceScanner scanner = runtime.openDatasetScanner(datasetLease, Optional.empty(), scanOptions)) { assertThat(readAllRows(scanner)).isEqualTo(3); } } @@ -112,13 +120,12 @@ public void testDatasetScanWithoutFragmentIdsClearsScanOptionsFragmentIds() .fragmentIds(List.of(fragments.get(0).getId())) .build(); - try (LanceScanner scanner = runtime.openDatasetScanner( + try (LanceRuntime.DatasetLease datasetLease = runtime.getDatasetLease( TestingConnectorSession.SESSION.getUser(), tableHandle.getTablePath(), tableHandle.getDatasetVersion(), - Optional.empty(), - scanOptions, - Collections.emptyMap())) { + Collections.emptyMap()); + LanceScanner scanner = runtime.openDatasetScanner(datasetLease, Optional.empty(), scanOptions)) { assertThat(readAllRows(scanner)).isEqualTo(4); } } @@ -150,13 +157,13 @@ public void testAllFragmentsSplitAppliesFilterAndLimitThroughPageSource() filteredLimitHandle, List.of(colX), null)) { - Page page = pageSource.getNextPage(); + Page page = pageSource.getNextSourcePage().getPage(); assertThat(page).isNotNull(); assertThat(page.getChannelCount()).isEqualTo(1); assertThat(page.getPositionCount()).isEqualTo(1); assertThat(BIGINT.getLong(page.getBlock(0), 0)).isGreaterThanOrEqualTo(2L); - assertThat(pageSource.getNextPage()).isNull(); + assertThat(pageSource.getNextSourcePage()).isNull(); assertThat(pageSource.isFinished()).isTrue(); } } @@ -174,7 +181,7 @@ public void testFragmentScan() List columns = runtime.getColumnHandleList(null, lanceTableHandle.getTablePath(), null, Collections.emptyMap()); // testing split 0 is enough try (LanceFragmentPageSource pageSource = new LanceFragmentPageSource(lanceTableHandle, columns, lanceSplit.getFragments(), Collections.emptyMap(), 8192, null, runtime)) { - Page page = pageSource.getNextPage(); + Page page = pageSource.getNextSourcePage().getPage(); // assert row/column count assertThat(page.getChannelCount()).isEqualTo(4); assertThat(page.getPositionCount()).isEqualTo(2); @@ -184,8 +191,7 @@ public void testFragmentScan() block = page.getBlock(1); assertThat(BIGINT.getLong(block, 1)).isEqualTo(2L); // assert no second page. it should come from the other split - page = pageSource.getNextPage(); - assertThat(page).isNull(); + assertThat(pageSource.getNextSourcePage()).isNull(); // assert that page is now finish assertThat(pageSource.isFinished()).isTrue(); } @@ -220,7 +226,7 @@ public void testColumnProjection() 8192, null, runtime)) { - Page page = pageSource.getNextPage(); + Page page = pageSource.getNextSourcePage().getPage(); assertThat(page.getChannelCount()).isEqualTo(2); assertThat(page.getPositionCount()).isEqualTo(2); @@ -264,7 +270,7 @@ public void testPartialColumnProjection() 8192, null, runtime)) { - Page page = pageSource.getNextPage(); + Page page = pageSource.getNextSourcePage().getPage(); // assert only 2 columns returned assertThat(page.getChannelCount()).isEqualTo(2); diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceMetadata.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceMetadata.java index 342e391..1f61950 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceMetadata.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceMetadata.java @@ -18,11 +18,14 @@ import com.google.common.collect.ImmutableSet; import com.google.common.io.Resources; import io.airlift.json.JsonCodec; +import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.connector.RetryMode; import io.trino.spi.connector.SchemaTableName; import io.trino.spi.connector.TableNotFoundException; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.FieldType; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; @@ -42,10 +45,18 @@ public class TestLanceMetadata { // Use URL.toString() to match the format used by LanceNamespaceHolder (file:/... vs file:///...) private static final String TEST_DB_PATH = Resources.getResource(TestLanceMetadata.class, "/example_db").toString() + "/"; - private static final LanceTableHandle TEST_TABLE_1_HANDLE = new LanceTableHandle("default", "test_table1", - TEST_DB_PATH + "test_table1.lance", List.of("test_table1"), Map.of()); - private static final LanceTableHandle TEST_TABLE_2_HANDLE = new LanceTableHandle("default", "test_table2", - TEST_DB_PATH + "test_table2.lance", List.of("test_table2"), Map.of()); + private static final LanceTableHandle TEST_TABLE_1_HANDLE = new LanceTableHandle( + "default", + "test_table1", + TEST_DB_PATH + "test_table1.lance", + List.of("test_table1"), + Map.of()); + private static final LanceTableHandle TEST_TABLE_2_HANDLE = new LanceTableHandle( + "default", + "test_table2", + TEST_DB_PATH + "test_table2.lance", + List.of("test_table2"), + Map.of()); // Actual column order in test data: x, y, b, c (field IDs 0, 1, 2, 3) private static final ArrowType INT64_TYPE = new ArrowType.Int(64, true); @@ -69,6 +80,14 @@ public void setUp() metadata = new LanceMetadata(runtime, lanceConfig, commitTaskDataCodec, mergeCommitDataCodec); } + @AfterEach + public void tearDown() + { + if (runtime != null) { + runtime.close(); + } + } + @Test public void testListSchemaNames() { @@ -153,6 +172,51 @@ public void testBuildPositionalOrdinalsWithNonSequentialFieldIds() .containsEntry("e", 3); } + @Test + public void testDatasetLeaseSurvivesCacheInvalidation() + { + LanceTableHandle table = metadata.getTableHandle(SESSION, new SchemaTableName("default", "test_table1"), Optional.empty(), Optional.empty()); + + try (LanceRuntime.DatasetLease datasetLease = runtime.getDatasetLease( + SESSION.getUser(), + table.getTablePath(), + table.getDatasetVersion(), + table.getStorageOptions())) { + runtime.invalidate(SESSION.getUser(), table.getTablePath()); + + assertThat(datasetLease.getDataset().getSchema().getFields()).isNotEmpty(); + } + } + + @Test + public void testExpiringStorageOptionsBypassDatasetCache() + { + LanceTableHandle table = metadata.getTableHandle(SESSION, new SchemaTableName("default", "test_table1"), Optional.empty(), Optional.empty()); + Map expiringStorageOptions = ImmutableMap.of( + "expires_at_millis", Long.toString(System.currentTimeMillis() + 60 * 60 * 1000)); + + assertThat(runtime.getFragments( + SESSION.getUser(), + table.getTablePath(), + table.getDatasetVersion(), + expiringStorageOptions)) + .isNotEmpty(); + assertThat(runtime.getCachedDatasetCount()).isZero(); + } + + @Test + public void testCleanupQueryClosesAbandonedTransactionDatasets() + { + LanceTableHandle table = metadata.getTableHandle(SESSION, new SchemaTableName("default", "test_table1"), Optional.empty(), Optional.empty()); + List columns = List.copyOf(metadata.getColumnHandles(SESSION, table).values()); + + metadata.beginInsert(SESSION, table, columns, RetryMode.NO_RETRIES); + + assertThat(metadata.getTransactionDatasetCount()).isEqualTo(1); + metadata.cleanupQuery(SESSION); + assertThat(metadata.getTransactionDatasetCount()).isZero(); + } + @Test public void testListTables() { diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceSplitManager.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceSplitManager.java index 68712b6..89ec72d 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceSplitManager.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceSplitManager.java @@ -20,6 +20,7 @@ import io.trino.spi.connector.ConnectorTableHandle; import io.trino.spi.connector.SchemaTableName; import io.trino.testing.TestingConnectorSession; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInstance; @@ -63,6 +64,14 @@ public void setUp() this.splitManager = new LanceSplitManager(runtime); } + @AfterEach + public void tearDown() + { + if (runtime != null) { + runtime.close(); + } + } + @Test public void testFullScanCreatesSplitPerFragment() throws ExecutionException, InterruptedException @@ -198,7 +207,6 @@ public void testAllFragmentsSplit() assertThat(split.isAllFragments()).isTrue(); assertThat(split.getFragments()).isEmpty(); - assertThat(split.getSplitInfo()).containsEntry("fragments", "ALL"); JsonCodec codec = JsonCodec.jsonCodec(LanceSplit.class); String json = codec.toJson(split); diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceStructColumns.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceStructColumns.java index 9a917f5..db6c8d5 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceStructColumns.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceStructColumns.java @@ -109,9 +109,10 @@ public void testDescribeTableWithStructColumn() assertUpdate("CREATE TABLE " + tableName + " (id BIGINT, metadata ROW(name VARCHAR, value BIGINT))"); - assertQuery("SELECT column_name, data_type FROM information_schema.columns " + + assertQuery( + "SELECT column_name, data_type FROM information_schema.columns " + "WHERE table_name = '" + tableName + "' ORDER BY ordinal_position", - "VALUES ('id', 'bigint'), ('metadata', 'row(name varchar, value bigint)')"); + "VALUES ('id', 'bigint'), ('metadata', 'row(\"name\" varchar, \"value\" bigint)')"); } finally { assertUpdate("DROP TABLE IF EXISTS " + tableName); diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceVectorColumns.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceVectorColumns.java index 1d303b9..45b8104 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceVectorColumns.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceVectorColumns.java @@ -57,7 +57,8 @@ public void testCreateTableAsSelectWithVectorColumn() String tableName = "test_vector_ctas_" + System.currentTimeMillis(); try { // Create table with vector column via CTAS - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (vector_columns = 'embedding:3') AS " + "SELECT CAST(1 AS BIGINT) as id, ARRAY[1.0E0, 2.0E0, 3.0E0] as embedding", 1); @@ -77,7 +78,8 @@ public void testInsertIntoVectorColumn() String tableName = "test_vector_insert_" + System.currentTimeMillis(); try { // Create table with vector column - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (vector_columns = 'embedding:3') AS " + "SELECT CAST(1 AS BIGINT) as id, ARRAY[1.0E0, 2.0E0, 3.0E0] as embedding", 1); @@ -121,7 +123,8 @@ public void testMultipleVectorColumns() String tableName = "test_multi_vector_" + System.currentTimeMillis(); try { // Create table with multiple vector columns - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (vector_columns = 'embedding1:2, embedding2:3') AS " + "SELECT CAST(1 AS BIGINT) as id, " + " ARRAY[1.0E0, 2.0E0] as embedding1, " + @@ -142,7 +145,8 @@ public void testVectorWithDoubleType() String tableName = "test_vector_double_" + System.currentTimeMillis(); try { // Create table with ARRAY(DOUBLE) vector column - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (vector_columns = 'embedding:3') AS " + "SELECT CAST(1 AS BIGINT) as id, " + " CAST(ARRAY[1.0E0, 2.0E0, 3.0E0] AS ARRAY(DOUBLE)) as embedding", @@ -162,7 +166,8 @@ public void testVectorWithMixedColumns() String tableName = "test_vector_mixed_" + System.currentTimeMillis(); try { // Create table with mix of vector and non-vector columns - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (vector_columns = 'vector_col:3') AS " + "SELECT CAST(1 AS BIGINT) as id, " + " 'regular text' as text_content, " + @@ -185,7 +190,8 @@ public void testBlobAndVectorCombined() String tableName = "test_blob_vector_combined_" + System.currentTimeMillis(); try { // Create table with both blob and vector columns - assertUpdate("CREATE TABLE " + tableName + " " + + assertUpdate( + "CREATE TABLE " + tableName + " " + "WITH (blob_columns = 'content', vector_columns = 'embedding:3') AS " + "SELECT CAST(1 AS BIGINT) as id, " + " X'48454C4C4F' as content, " + diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestSubstraitExpressionBuilder.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestSubstraitExpressionBuilder.java index f54dd24..a6a94f8 100644 --- a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestSubstraitExpressionBuilder.java +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestSubstraitExpressionBuilder.java @@ -78,7 +78,7 @@ public void testNoneDomain() Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); // The result should be a false literal - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -88,7 +88,7 @@ public void testSingleEquality() Map.of(INT_COLUMN, Domain.singleValue(INTEGER, 42L))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -98,7 +98,7 @@ public void testVarcharEquality() Map.of(VARCHAR_COLUMN, Domain.singleValue(VARCHAR, Slices.utf8Slice("test")))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -108,7 +108,7 @@ public void testBooleanValue() Map.of(BOOLEAN_COLUMN, Domain.singleValue(BOOLEAN, true))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -118,7 +118,7 @@ public void testInClause() Map.of(INT_COLUMN, Domain.multipleValues(INTEGER, java.util.List.of(1L, 2L, 3L)))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -129,7 +129,7 @@ public void testRangeGreaterThan() ValueSet.ofRanges(Range.greaterThan(INTEGER, 10L)), false))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -140,7 +140,7 @@ public void testRangeBetween() ValueSet.ofRanges(Range.range(INTEGER, 10L, true, 100L, true)), false))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -150,7 +150,7 @@ public void testIsNull() Map.of(INT_COLUMN, Domain.onlyNull(INTEGER))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -160,7 +160,7 @@ public void testIsNotNull() Map.of(INT_COLUMN, Domain.notNull(INTEGER))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -172,7 +172,7 @@ public void testMultipleColumns() VARCHAR_COLUMN, Domain.singleValue(VARCHAR, Slices.utf8Slice("test")))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -183,7 +183,7 @@ public void testExpressionConversionDirect() Optional result = SubstraitExpressionBuilder.tupleDomainToExpression(domain, COLUMN_ORDINALS); assertThat(result).isPresent(); // Verify it's a scalar function invocation (equality) - assertThat(result.get()).isInstanceOf(Expression.ScalarFunctionInvocation.class); + assertThat(result.orElseThrow()).isInstanceOf(Expression.ScalarFunctionInvocation.class); } @Test @@ -220,7 +220,7 @@ public void testTimestampMicrosEquality() Map.of(tsColumn, Domain.singleValue(TIMESTAMP_MICROS, epochMicros))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, columns, ordinals); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -237,7 +237,7 @@ public void testTimestampMillisEquality() Map.of(tsColumn, Domain.singleValue(TIMESTAMP_MILLIS, epochMicros))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, columns, ordinals); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -254,7 +254,7 @@ public void testTimestampRange() ValueSet.ofRanges(Range.range(TIMESTAMP_MICROS, startMicros, true, endMicros, false)), false))); Optional result = SubstraitExpressionBuilder.tupleDomainToSubstrait(domain, columns, ordinals); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test @@ -378,7 +378,7 @@ public void testCombineExpressionsWithLike() tupleDomainExpr, likePredicates, ALL_COLUMNS, COLUMN_ORDINALS); assertThat(result).isPresent(); - assertThat(result.get().remaining()).isGreaterThan(0); + assertThat(result.orElseThrow().remaining()).isGreaterThan(0); } @Test diff --git a/pom.xml b/pom.xml index e0ac0e7..359faae 100644 --- a/pom.xml +++ b/pom.xml @@ -5,7 +5,7 @@ io.airlift airbase - 264 + 375 org.lance @@ -55,26 +55,30 @@ - 476 + 481 - 23 + 25 true true true - 23.0.0 - 8 + 25.0.1 + ${project.build.targetJdk} 3g -XX:G1HeapRegionSize=32M -XX:+UnlockDiagnosticVMOptions - -XX:+EnableDynamicAgentLoading ${extraJavaVectorArgs} - -Djava.security.manager=allow + --sun-misc-unsafe-memory-access=allow + -Dairlift.quiet=true ${air.test.jvm.additional-arguments.default} - 336 + 424 4.13.2 - 3.9.0 - 3.25.5 + 1.12.1 + 1.76.3 + 4.2.0 + 4.2.12.Final + 4.34.1 + 2.0.77.Final --add-modules=jdk.incubator.vector GA @@ -90,10 +94,18 @@ import + + io.grpc + grpc-bom + ${dep.grpc.version} + pom + import + + io.netty netty-bom - 4.1.119.Final + ${dep.netty.version} pom import @@ -101,7 +113,7 @@ org.apache.arrow arrow-bom - 18.3.0 + 19.0.0 pom import @@ -109,7 +121,7 @@ org.testcontainers testcontainers-bom - 1.20.4 + 2.0.5 pom import @@ -117,7 +129,7 @@ software.amazon.awssdk bom - 2.29.29 + 2.44.0 pom import @@ -134,6 +146,88 @@ 2 + + io.netty + netty-tcnative + ${dep.tcnative.version} + linux-x86_64 + + + + io.netty + netty-tcnative + ${dep.tcnative.version} + linux-x86_64-fedora + + + + io.netty + netty-tcnative + ${dep.tcnative.version} + linux-aarch_64-fedora + + + + io.netty + netty-tcnative + ${dep.tcnative.version} + osx-aarch_64 + + + + io.netty + netty-tcnative + ${dep.tcnative.version} + osx-x86_64 + + + + io.netty + netty-tcnative-boringssl-static + ${dep.tcnative.version} + + + + io.netty + netty-tcnative-boringssl-static + ${dep.tcnative.version} + linux-x86_64 + + + + io.netty + netty-tcnative-boringssl-static + ${dep.tcnative.version} + linux-aarch_64 + + + + io.netty + netty-tcnative-boringssl-static + ${dep.tcnative.version} + osx-x86_64 + + + + io.netty + netty-tcnative-boringssl-static + ${dep.tcnative.version} + osx-aarch_64 + + + + io.netty + netty-tcnative-boringssl-static + ${dep.tcnative.version} + windows-x86_64 + + + + io.netty + netty-tcnative-classes + ${dep.tcnative.version} + + io.trino trino-cache @@ -172,25 +266,12 @@ ${trino.version} - - io.trino - trino-spi - ${trino.version} - test-jar - - io.trino trino-testing ${trino.version} - - io.trino - trino-testing-kafka - ${trino.version} - - io.trino trino-tpch @@ -212,13 +293,13 @@ org.apache.avro avro - 1.12.0 + ${dep.avro.version} org.apache.commons commons-lang3 - 3.17.0 + 3.20.0 @@ -230,14 +311,9 @@ org.jetbrains annotations - 26.0.1 + 26.1.0 - - org.openjdk.jol - jol-core - 0.17 - @@ -264,6 +340,25 @@ mozilla/public-suffix-list.txt + + + + org.alluxio + alluxio-core-client-fs + + + org.alluxio + alluxio-core-common + + + org.alluxio + alluxio-core-transport + + + + git.properties + + @@ -310,7 +405,7 @@ io.trino trino-maven-plugin - 15 + 20 true