diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCompactHandle.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCompactHandle.java new file mode 100644 index 0000000..d037448 --- /dev/null +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCompactHandle.java @@ -0,0 +1,33 @@ +/* + * 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 java.util.Optional; + +import static java.util.Objects.requireNonNull; + +/** + * Procedure handle for {@code ALTER TABLE ... EXECUTE compact(...)}. + * {@code deferIndexRemap}, when true, skips remapping row addresses in existing indices during + * compaction (applied lazily via the Fragment Reuse Index on next load instead), so compaction + * does not conflict with concurrent index builds. Empty means use Lance's own default. + */ +public record LanceCompactHandle(Optional deferIndexRemap) + implements LanceProcedureHandle +{ + public LanceCompactHandle + { + requireNonNull(deferIndexRemap, "deferIndexRemap is null"); + } +} 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..f279e98 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 @@ -22,10 +22,12 @@ import io.trino.spi.connector.ConnectorSession; import io.trino.spi.connector.ConnectorSplitManager; import io.trino.spi.connector.ConnectorTransactionHandle; +import io.trino.spi.connector.TableProcedureMetadata; import io.trino.spi.session.PropertyMetadata; import io.trino.spi.transaction.IsolationLevel; import java.util.List; +import java.util.Set; import static java.util.Objects.requireNonNull; @@ -90,6 +92,12 @@ public List> getTableProperties() return LanceTableProperties.getTableProperties(); } + @Override + public Set getTableProcedures() + { + return LanceTableProcedures.getTableProcedures(); + } + @Override public final void shutdown() { diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCreateIndexHandle.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCreateIndexHandle.java new file mode 100644 index 0000000..72ca92e --- /dev/null +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceCreateIndexHandle.java @@ -0,0 +1,49 @@ +/* + * 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 java.util.Optional; + +import static java.util.Objects.requireNonNull; + +/** + * Procedure handle for {@code ALTER TABLE ... EXECUTE create_index(...)}. + * Currently only the {@code fts} (inverted) index type is supported. + */ +public record LanceCreateIndexHandle( + String column, + Optional indexName, + String indexType, + boolean replace, + boolean train, + String baseTokenizer, + String language, + boolean withPosition, + boolean lowerCase, + boolean stem, + boolean removeStopWords, + boolean asciiFolding, + Optional maxTokenLength) + implements LanceProcedureHandle +{ + public LanceCreateIndexHandle + { + requireNonNull(column, "column is null"); + requireNonNull(indexName, "indexName is null"); + requireNonNull(indexType, "indexType is null"); + requireNonNull(baseTokenizer, "baseTokenizer is null"); + requireNonNull(language, "language is null"); + requireNonNull(maxTokenLength, "maxTokenLength is null"); + } +} 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..0168125 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 @@ -16,6 +16,7 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Lists; import com.google.inject.Inject; import io.airlift.json.JsonCodec; import io.airlift.log.Logger; @@ -26,12 +27,14 @@ import io.trino.spi.connector.Assignment; import io.trino.spi.connector.ColumnHandle; import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorAccessControl; import io.trino.spi.connector.ConnectorInsertTableHandle; import io.trino.spi.connector.ConnectorMergeTableHandle; import io.trino.spi.connector.ConnectorMetadata; import io.trino.spi.connector.ConnectorOutputMetadata; import io.trino.spi.connector.ConnectorOutputTableHandle; import io.trino.spi.connector.ConnectorSession; +import io.trino.spi.connector.ConnectorTableExecuteHandle; import io.trino.spi.connector.ConnectorTableHandle; import io.trino.spi.connector.ConnectorTableLayout; import io.trino.spi.connector.ConnectorTableMetadata; @@ -65,11 +68,20 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.lance.CommitBuilder; import org.lance.Dataset; +import org.lance.Fragment; import org.lance.FragmentMetadata; import org.lance.ManifestSummary; import org.lance.ReadOptions; import org.lance.SourcedTransaction; import org.lance.Transaction; +import org.lance.compaction.CompactionOptions; +import org.lance.index.Index; +import org.lance.index.IndexOptions; +import org.lance.index.IndexParams; +import org.lance.index.IndexType; +import org.lance.index.OptimizeOptions; +import org.lance.index.scalar.InvertedIndexParams; +import org.lance.index.scalar.ScalarIndexParams; import org.lance.namespace.LanceNamespace; import org.lance.namespace.model.CreateNamespaceRequest; import org.lance.namespace.model.DeclareTableRequest; @@ -114,6 +126,12 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import static com.google.common.collect.ImmutableList.toImmutableList; @@ -1270,6 +1288,260 @@ public void finishMerge( } } + // ===== Table Procedures (ALTER TABLE ... EXECUTE) ===== + + @Override + public Optional getTableHandleForExecute( + ConnectorSession session, + ConnectorAccessControl accessControl, + ConnectorTableHandle tableHandle, + String procedureName, + Map executeProperties, + RetryMode retryMode) + { + LanceTableHandle table = (LanceTableHandle) tableHandle; + + LanceTableProcedureId procedureId; + try { + procedureId = LanceTableProcedureId.valueOf(procedureName); + } + catch (IllegalArgumentException e) { + throw new TrinoException(NOT_SUPPORTED, "Unknown table procedure: " + procedureName); + } + + LanceProcedureHandle procedureHandle = switch (procedureId) { + case CREATE_INDEX -> buildCreateIndexHandle(executeProperties); + case OPTIMIZE_INDICES -> buildOptimizeIndicesHandle(executeProperties); + case COMPACT -> buildCompactHandle(executeProperties); + }; + + Map storageOptions = getEffectiveStorageOptions(table); + log.debug("getTableHandleForExecute: table=%s, procedure=%s, handle=%s", table.getTableName(), procedureId, procedureHandle); + + return Optional.of(new LanceTableExecuteHandle( + new SchemaTableName(table.getSchemaName(), table.getTableName()), + table.getTablePath(), + table.getTableId(), + storageOptions, + procedureId, + procedureHandle)); + } + + private static LanceCreateIndexHandle buildCreateIndexHandle(Map properties) + { + String column = (String) properties.get(LanceTableProcedures.COLUMN); + if (column == null || column.isEmpty()) { + throw new TrinoException(INVALID_ARGUMENTS, "create_index requires a 'column' argument"); + } + String indexType = (String) properties.get(LanceTableProcedures.INDEX_TYPE); + if (!"fts".equalsIgnoreCase(indexType)) { + throw new TrinoException(NOT_SUPPORTED, "create_index currently only supports index_type => 'fts', got: " + indexType); + } + + return new LanceCreateIndexHandle( + column, + Optional.ofNullable((String) properties.get(LanceTableProcedures.INDEX_NAME)), + indexType, + (Boolean) properties.get(LanceTableProcedures.REPLACE), + (Boolean) properties.get(LanceTableProcedures.TRAIN), + (String) properties.get(LanceTableProcedures.BASE_TOKENIZER), + (String) properties.get(LanceTableProcedures.LANGUAGE), + (Boolean) properties.get(LanceTableProcedures.WITH_POSITION), + (Boolean) properties.get(LanceTableProcedures.LOWER_CASE), + (Boolean) properties.get(LanceTableProcedures.STEM), + (Boolean) properties.get(LanceTableProcedures.REMOVE_STOP_WORDS), + (Boolean) properties.get(LanceTableProcedures.ASCII_FOLDING), + Optional.ofNullable((Integer) properties.get(LanceTableProcedures.MAX_TOKEN_LENGTH))); + } + + private static LanceOptimizeIndicesHandle buildOptimizeIndicesHandle(Map properties) + { + String indexNamesCsv = (String) properties.get(LanceTableProcedures.INDEX_NAMES); + List indexNames = List.of(); + if (indexNamesCsv != null && !indexNamesCsv.isEmpty()) { + indexNames = Arrays.stream(indexNamesCsv.split(",")) + .map(String::trim) + .filter(name -> !name.isEmpty()) + .collect(toImmutableList()); + } + + return new LanceOptimizeIndicesHandle( + indexNames, + Optional.ofNullable((Integer) properties.get(LanceTableProcedures.NUM_INDICES_TO_MERGE)), + (Boolean) properties.get(LanceTableProcedures.RETRAIN)); + } + + private static LanceCompactHandle buildCompactHandle(Map properties) + { + return new LanceCompactHandle(Optional.ofNullable((Boolean) properties.get(LanceTableProcedures.DEFER_INDEX_REMAP))); + } + + @Override + public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteHandle tableExecuteHandle) + { + LanceTableExecuteHandle handle = (LanceTableExecuteHandle) tableExecuteHandle; + String userIdentity = session.getUser(); + + log.debug("executeTableExecute: table=%s, procedure=%s, handle=%s", + handle.schemaTableName(), handle.procedureId(), handle.procedureHandle()); + + try (Dataset dataset = runtime.openDatasetDirect(userIdentity, handle.tablePath(), null, handle.storageOptions())) { + switch (handle.procedureId()) { + case CREATE_INDEX -> executeCreateIndex(dataset, handle.tablePath(), handle.storageOptions(), (LanceCreateIndexHandle) handle.procedureHandle()); + case OPTIMIZE_INDICES -> executeOptimizeIndices(dataset, (LanceOptimizeIndicesHandle) handle.procedureHandle()); + case COMPACT -> executeCompact(dataset, (LanceCompactHandle) handle.procedureHandle()); + } + } + catch (RuntimeException e) { + if (isCommitConflict(e)) { + throw new TrinoException(TRANSACTION_CONFLICT, "Concurrent modification conflict", e); + } + throw e; + } + + runtime.invalidate(userIdentity, handle.tablePath()); + } + + private void executeCreateIndex(Dataset dataset, String tablePath, Map storageOptions, LanceCreateIndexHandle handle) + { + IndexParams indexParams = IndexParams.builder().setScalarIndexParams(buildInvertedIndexParams(handle)).build(); + + // Registering an empty index (train=false) or training on a single-fragment table isn't + // worth parallelizing; fewer than 2 fragments means there's nothing to split work across. + List fragmentIds = handle.train() ? dataset.getFragments().stream().map(Fragment::getId).toList() : List.of(); + if (fragmentIds.size() < 2) { + IndexOptions.Builder optionsBuilder = IndexOptions + .builder(List.of(handle.column()), IndexType.INVERTED, indexParams) + .replace(handle.replace()) + .train(handle.train()); + handle.indexName().ifPresent(optionsBuilder::withIndexName); + + Index index = dataset.createIndex(optionsBuilder.build()); + log.debug("executeCreateIndex: created index %s on column %s", index.name(), handle.column()); + return; + } + + // Phase 2: build one index segment per batch of fragments in parallel, consolidate the + // resulting segments into one physical segment, then commit that single segment as the named + // logical index in one write-locked call. This parallelizes the CPU-bound tokenize/build work + // across the coordinator's cores; it does not distribute across Trino worker nodes (see #187/ + // #188 for why: Trino's distributed table-execute machinery requires scanning full table rows + // through the page pipeline, which would force redundant reads on top of Lance's own native + // fragment scan). + // + // Each batch opens its OWN independent Dataset handle (see buildIndexSegmentsInParallel) rather + // than sharing this method's `dataset` - benchmarking showed concurrent createIndex calls + // against one shared Dataset instance are consistently slower than the plain single-call path + // (shared allocator/handle contention, not anything about Lance's own createIndex being + // internally parallelized - it isn't; lance-index's inverted-index builder has no internal + // rayon/tokio/thread concurrency), while independent per-thread handles show real ~20-40% + // speedup up to about 4 concurrent builds before plateauing. + // + // Each batch is built with its own auto-generated UUID (IndexOptions.withIndexUUID is left + // unset): mergeExistingIndexSegments rejects segments sharing one UUID as duplicates, so - despite + // its Javadoc ("multiple fragment-level indices need to share UUID for later merging") - passing + // one shared UUID across batches does not work; the segments passed to mergeExistingIndexSegments + // must each carry a distinct identity. + String indexName = handle.indexName().orElseGet(() -> handle.column() + "_idx"); + int parallelism = Math.min(fragmentIds.size(), Runtime.getRuntime().availableProcessors()); + + log.info("executeCreateIndex: building index %s on column %s across %d fragments with parallelism %d", + indexName, handle.column(), fragmentIds.size(), parallelism); + + // Each per-fragment createIndex call validates the index name against the dataset's currently + // committed indices up front, even though it doesn't commit anything itself - so an existing + // index of the same name must be dropped before building segments, not just before the final + // commit, or every batch fails with "Index name '...' already exists". + if (handle.replace() && dataset.listIndexes().contains(indexName)) { + dataset.dropIndex(indexName); + } + + List segments = buildIndexSegmentsInParallel(tablePath, storageOptions, handle.column(), indexParams, fragmentIds, parallelism); + Index mergedSegment = dataset.mergeExistingIndexSegments(segments); + List committed = dataset.commitExistingIndexSegments(indexName, handle.column(), List.of(mergedSegment)); + log.debug("executeCreateIndex: committed %d fragment-parallel segments as index %s", committed.size(), indexName); + } + + @VisibleForTesting + static ScalarIndexParams buildInvertedIndexParams(LanceCreateIndexHandle handle) + { + InvertedIndexParams.Builder tokenizerBuilder = InvertedIndexParams.builder() + .baseTokenizer(handle.baseTokenizer()) + .language(handle.language()) + .withPosition(handle.withPosition()) + .lowerCase(handle.lowerCase()) + .stem(handle.stem()) + .removeStopWords(handle.removeStopWords()) + .asciiFolding(handle.asciiFolding()); + handle.maxTokenLength().ifPresent(tokenizerBuilder::maxTokenLength); + return tokenizerBuilder.build(); + } + + @VisibleForTesting + static List buildIndexSegmentsInParallel( + String tablePath, Map storageOptions, String column, IndexParams indexParams, List fragmentIds, int parallelism) + { + int batchSize = (fragmentIds.size() + parallelism - 1) / parallelism; + List> batches = Lists.partition(fragmentIds, batchSize); + AtomicInteger threadCounter = new AtomicInteger(); + ThreadFactory threadFactory = runnable -> { + Thread thread = new Thread(runnable, "lance-create-index-" + threadCounter.getAndIncrement()); + thread.setDaemon(true); + return thread; + }; + try (ExecutorService executor = Executors.newFixedThreadPool(batches.size(), threadFactory)) { + // Each batch opens its own independent Dataset handle (own native handle, own allocator - + // not a shared session/allocator via LanceRuntime) and gets its own auto-generated segment + // UUID (withIndexUUID left unset - see the caller for why a shared UUID across batches + // doesn't work here). + List> futures = batches.stream() + .map(batch -> executor.submit(() -> { + ReadOptions.Builder readOptionsBuilder = new ReadOptions.Builder(); + if (storageOptions != null && !storageOptions.isEmpty()) { + readOptionsBuilder.setStorageOptions(storageOptions); + } + try (Dataset batchDataset = Dataset.open(tablePath, readOptionsBuilder.build())) { + return batchDataset.createIndex( + IndexOptions.builder(List.of(column), IndexType.INVERTED, indexParams) + .withFragmentIds(batch) + .train(true) + .build()); + } + })) + .collect(toImmutableList()); + + List segments = new ArrayList<>(); + for (Future future : futures) { + segments.add(future.get()); + } + return segments; + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new TrinoException(GENERIC_INTERNAL_ERROR, "Interrupted while building index segments", e); + } + catch (ExecutionException e) { + throw new TrinoException(GENERIC_INTERNAL_ERROR, "Failed to build an index segment", e.getCause()); + } + } + + private void executeOptimizeIndices(Dataset dataset, LanceOptimizeIndicesHandle handle) + { + OptimizeOptions.Builder builder = OptimizeOptions.builder().retrain(handle.retrain()); + if (!handle.indexNames().isEmpty()) { + builder.indexNames(handle.indexNames()); + } + handle.numIndicesToMerge().ifPresent(builder::numIndicesToMerge); + dataset.optimizeIndices(builder.build()); + } + + private void executeCompact(Dataset dataset, LanceCompactHandle handle) + { + CompactionOptions.Builder builder = CompactionOptions.builder(); + handle.deferIndexRemap().ifPresent(builder::withDeferIndexRemap); + dataset.compact(builder.build()); + } + // ===== Helper Methods ===== private LanceNamespace getNamespace() diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceOptimizeIndicesHandle.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceOptimizeIndicesHandle.java new file mode 100644 index 0000000..4500cbb --- /dev/null +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceOptimizeIndicesHandle.java @@ -0,0 +1,34 @@ +/* + * 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 java.util.List; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +/** + * Procedure handle for {@code ALTER TABLE ... EXECUTE optimize_indices(...)}. + * Incrementally indexes fragments that were appended since an index was last built or optimized, + * without a full rebuild. {@code indexNames} empty means all indices on the table. + */ +public record LanceOptimizeIndicesHandle(List indexNames, Optional numIndicesToMerge, boolean retrain) + implements LanceProcedureHandle +{ + public LanceOptimizeIndicesHandle + { + indexNames = List.copyOf(requireNonNull(indexNames, "indexNames is null")); + requireNonNull(numIndicesToMerge, "numIndicesToMerge is null"); + } +} diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceProcedureHandle.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceProcedureHandle.java new file mode 100644 index 0000000..4779757 --- /dev/null +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceProcedureHandle.java @@ -0,0 +1,31 @@ +/* + * 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 com.fasterxml.jackson.annotation.JsonSubTypes; +import com.fasterxml.jackson.annotation.JsonTypeInfo; + +/** + * Marker interface for procedure-specific {@link LanceTableExecuteHandle} payloads. + * {@code LanceTableExecuteHandle} is serialized as part of the query plan sent from the + * coordinator to itself for execution, so the concrete subtype needs explicit type info for + * Jackson to deserialize it back correctly. + */ +@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "@type") +@JsonSubTypes({ + @JsonSubTypes.Type(value = LanceCreateIndexHandle.class, name = "create_index"), + @JsonSubTypes.Type(value = LanceOptimizeIndicesHandle.class, name = "optimize_indices"), + @JsonSubTypes.Type(value = LanceCompactHandle.class, name = "compact"), +}) +public interface LanceProcedureHandle {} diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableExecuteHandle.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableExecuteHandle.java new file mode 100644 index 0000000..ffaf4f8 --- /dev/null +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableExecuteHandle.java @@ -0,0 +1,53 @@ +/* + * 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 com.google.common.collect.ImmutableMap; +import io.trino.spi.connector.ConnectorTableExecuteHandle; +import io.trino.spi.connector.SchemaTableName; + +import java.util.List; +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +/** + * Handle for an in-flight {@code ALTER TABLE ... EXECUTE} table procedure on a Lance table. + * {@code procedureHandle} carries the procedure-specific arguments, e.g. {@link LanceCreateIndexHandle}. + */ +public record LanceTableExecuteHandle( + SchemaTableName schemaTableName, + String tablePath, + List tableId, + Map storageOptions, + LanceTableProcedureId procedureId, + LanceProcedureHandle procedureHandle) + implements ConnectorTableExecuteHandle +{ + public LanceTableExecuteHandle + { + requireNonNull(schemaTableName, "schemaTableName is null"); + requireNonNull(tablePath, "tablePath is null"); + tableId = List.copyOf(requireNonNull(tableId, "tableId is null")); + storageOptions = ImmutableMap.copyOf(requireNonNull(storageOptions, "storageOptions is null")); + requireNonNull(procedureId, "procedureId is null"); + requireNonNull(procedureHandle, "procedureHandle is null"); + } + + @Override + public String toString() + { + return "schemaTableName:%s, procedureId:%s, procedureHandle:{%s}".formatted(schemaTableName, procedureId, procedureHandle); + } +} diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProcedureId.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProcedureId.java new file mode 100644 index 0000000..f8fb48a --- /dev/null +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProcedureId.java @@ -0,0 +1,26 @@ +/* + * 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; + +/** + * Identifiers for {@code ALTER TABLE ... EXECUTE} table procedures supported by the Lance connector. + * All procedures are currently coordinator-only: they call directly into {@code org.lance.Dataset} + * index/compaction APIs rather than reading or writing table data through Trino's split/page pipeline. + */ +public enum LanceTableProcedureId +{ + CREATE_INDEX, + OPTIMIZE_INDICES, + COMPACT, +} diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProcedures.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProcedures.java new file mode 100644 index 0000000..201aaf1 --- /dev/null +++ b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceTableProcedures.java @@ -0,0 +1,109 @@ +/* + * 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 com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import io.trino.spi.connector.TableProcedureMetadata; + +import java.util.Set; + +import static io.trino.spi.connector.TableProcedureExecutionMode.coordinatorOnly; +import static io.trino.spi.session.PropertyMetadata.booleanProperty; +import static io.trino.spi.session.PropertyMetadata.integerProperty; +import static io.trino.spi.session.PropertyMetadata.stringProperty; + +/** + * Table procedures exposed via {@code ALTER TABLE ... EXECUTE}. All procedures here build or + * maintain Lance indexes directly through {@code org.lance.Dataset} (index build, incremental + * optimize, compaction) and never read or write table data through Trino's split/page pipeline, + * so they all run coordinator-only. + */ +public final class LanceTableProcedures +{ + // CREATE_INDEX properties + public static final String COLUMN = "column"; + public static final String INDEX_NAME = "index_name"; + public static final String INDEX_TYPE = "index_type"; + public static final String REPLACE = "replace"; + public static final String TRAIN = "train"; + public static final String BASE_TOKENIZER = "base_tokenizer"; + public static final String LANGUAGE = "language"; + public static final String WITH_POSITION = "with_position"; + public static final String LOWER_CASE = "lower_case"; + public static final String STEM = "stem"; + public static final String REMOVE_STOP_WORDS = "remove_stop_words"; + public static final String ASCII_FOLDING = "ascii_folding"; + public static final String MAX_TOKEN_LENGTH = "max_token_length"; + + // OPTIMIZE_INDICES properties + public static final String INDEX_NAMES = "index_names"; + public static final String NUM_INDICES_TO_MERGE = "num_indices_to_merge"; + public static final String RETRAIN = "retrain"; + + // COMPACT properties + public static final String DEFER_INDEX_REMAP = "defer_index_remap"; + + private LanceTableProcedures() {} + + public static Set getTableProcedures() + { + return ImmutableSet.of(createIndex(), optimizeIndices(), compact()); + } + + private static TableProcedureMetadata createIndex() + { + return new TableProcedureMetadata( + LanceTableProcedureId.CREATE_INDEX.name(), + coordinatorOnly(), + ImmutableList.of( + stringProperty(COLUMN, "Column to build the index on", null, false), + stringProperty(INDEX_NAME, "Name for the index; auto-generated if not provided", null, false), + stringProperty(INDEX_TYPE, "Index type to create; currently only 'fts' is supported", "fts", false), + booleanProperty(REPLACE, "Replace an existing index with the same name", false, false), + booleanProperty(TRAIN, "Train the index on existing data now; if false, registers an empty index to populate later via optimize_indices", true, false), + stringProperty(BASE_TOKENIZER, "FTS tokenizer: simple, whitespace, raw, ngram, icu, icu/split, lindera/*, jieba/*", "simple", false), + stringProperty(LANGUAGE, "Language used for stemming and stop words", "English", false), + booleanProperty(WITH_POSITION, "Store token positions to support phrase queries", false, false), + booleanProperty(LOWER_CASE, "Lower-case tokens", true, false), + booleanProperty(STEM, "Apply stemming", false, false), + booleanProperty(REMOVE_STOP_WORDS, "Remove stop words", false, false), + booleanProperty(ASCII_FOLDING, "Apply ASCII folding", false, false), + integerProperty(MAX_TOKEN_LENGTH, "Maximum token length", null, false))); + } + + private static TableProcedureMetadata optimizeIndices() + { + return new TableProcedureMetadata( + LanceTableProcedureId.OPTIMIZE_INDICES.name(), + coordinatorOnly(), + ImmutableList.of( + stringProperty(INDEX_NAMES, "Comma-separated index names to optimize; all indices on the table if not specified", null, false), + integerProperty(NUM_INDICES_TO_MERGE, "Number of index segments to merge while optimizing", null, false), + booleanProperty(RETRAIN, "Retrain the index from scratch instead of incrementally indexing newly appended fragments", false, false))); + } + + private static TableProcedureMetadata compact() + { + return new TableProcedureMetadata( + LanceTableProcedureId.COMPACT.name(), + coordinatorOnly(), + ImmutableList.of( + booleanProperty( + DEFER_INDEX_REMAP, + "Defer remapping row addresses in indices to next index load, so compaction does not conflict with concurrent index builds", + null, + false))); + } +} diff --git a/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceTableExecuteProcedures.java b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceTableExecuteProcedures.java new file mode 100644 index 0000000..d26662d --- /dev/null +++ b/plugin/trino-lance/src/test/java/io/trino/plugin/lance/TestLanceTableExecuteProcedures.java @@ -0,0 +1,229 @@ +/* + * 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 io.trino.testing.AbstractTestQueryFramework; +import io.trino.testing.QueryRunner; +import org.junit.jupiter.api.Test; +import org.lance.Dataset; +import org.lance.Fragment; +import org.lance.ReadOptions; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Coverage for the {@code create_index} / {@code optimize_indices} / {@code compact} table + * procedures (Phase 1, coordinator-only). These call directly into {@code org.lance.Dataset} + * rather than Trino's split/page pipeline, so correctness is verified both through SQL and by + * opening the underlying Lance dataset directly to inspect index state. + */ +public class TestLanceTableExecuteProcedures + extends AbstractTestQueryFramework +{ + private Path tempDir; + + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + tempDir = Files.createTempDirectory("lance-trino-execute-test"); + tempDir.toFile().deleteOnExit(); + return LanceQueryRunner.builder() + .addConnectorProperty("lance.root", tempDir.toUri().toString()) + .addConnectorProperty("lance.single_level_ns", "true") + .build(); + } + + private Dataset openDataset(String tableName) + { + String tablePath = tempDir.toUri() + tableName + ".lance"; + return Dataset.open(tablePath, new ReadOptions.Builder().build()); + } + + @Test + public void testCreateIndexBuildsFtsIndex() + { + String tableName = "test_create_index_" + System.currentTimeMillis(); + try { + assertUpdate("CREATE TABLE " + tableName + " (id bigint, body varchar)"); + assertUpdate("INSERT INTO " + tableName + " VALUES (1, 'hello world'), (2, 'goodbye world')", 2); + + getQueryRunner().execute("ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'fts')"); + + try (Dataset dataset = openDataset(tableName)) { + List indexNames = dataset.listIndexes(); + assertThat(indexNames) + .as("an index should have been created on 'body'") + .isNotEmpty(); + } + + // Re-running create_index for the same column without replace => true should fail: + // the auto-generated index name collides with the one just created. + assertThatThrownBy(() -> getQueryRunner().execute( + "ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'fts')")); + + // With replace => true it should succeed. + getQueryRunner().execute( + "ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'fts', replace => true)"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + } + } + + @Test + public void testCreateIndexParallelizesAcrossMultipleFragments() + { + String tableName = "test_create_index_parallel_" + System.currentTimeMillis(); + try { + assertUpdate("CREATE TABLE " + tableName + " (id bigint, body varchar)"); + // Separate INSERT statements so the table has multiple fragments before create_index + // runs, exercising the fragment-parallel build (org.lance.Dataset#createIndex with + // withFragmentIds/withIndexUUID per batch) + commitExistingIndexSegments merge path, + // instead of the single-fragment path that skips parallelization. + assertUpdate("INSERT INTO " + tableName + " VALUES (1, 'alpha document')", 1); + assertUpdate("INSERT INTO " + tableName + " VALUES (2, 'beta document')", 1); + assertUpdate("INSERT INTO " + tableName + " VALUES (3, 'gamma document')", 1); + assertUpdate("INSERT INTO " + tableName + " VALUES (4, 'delta document')", 1); + + List fragmentIdsBeforeIndex; + try (Dataset dataset = openDataset(tableName)) { + fragmentIdsBeforeIndex = dataset.getFragments().stream().map(Fragment::getId).toList(); + } + assertThat(fragmentIdsBeforeIndex.size()) + .as("test setup should produce multiple fragments so the parallel path is exercised") + .isGreaterThan(1); + + getQueryRunner().execute("ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'fts')"); + + try (Dataset dataset = openDataset(tableName)) { + assertThat(dataset.listIndexes()).contains("body_idx"); + } + + // The merge (commitExistingIndexSegments) must not have lost or duplicated any rows. + assertQuery("SELECT count(*) FROM " + tableName, "SELECT 4"); + assertQuery("SELECT id FROM " + tableName + " ORDER BY id", "VALUES 1, 2, 3, 4"); + + // replace => true must go through the same drop-then-commit path as the single-fragment case. + getQueryRunner().execute( + "ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'fts', replace => true)"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + } + } + + @Test + public void testCreateIndexWithLanguageAgnosticTokenizers() + { + String tableName = "test_create_index_lang_agnostic_" + System.currentTimeMillis(); + try { + assertUpdate("CREATE TABLE " + tableName + " (id bigint, body varchar)"); + // Mixed-script content: whitespace-delimited English, plus Chinese and Japanese text + // with no whitespace between words at all. base_tokenizer => 'simple' (the default) + // would treat each CJK line as one giant token since it only splits on whitespace/punctuation. + assertUpdate("INSERT INTO " + tableName + " VALUES " + + "(1, '这是一个测试文档，用于验证语言无关的分词器'), " + + "(2, 'これはテストです、言語に依存しないトークナイザーを検証するためのものです'), " + + "(3, 'hello world, this is an english test document')", + 3); + + // 'ngram' is character-level and works uniformly regardless of script; the most + // conservative choice when the corpus mixes arbitrary/unknown scripts. Note: 'icu' - + // Unicode's own generic word-boundary segmentation - is documented on the lance-format/lance + // main branch but is NOT accepted by the lance-core:7.0.0 native library this connector + // pins (fails with "unknown base tokenizer icu"); it isn't available until a newer release. + getQueryRunner().execute( + "ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'fts', base_tokenizer => 'ngram')"); + try (Dataset dataset = openDataset(tableName)) { + assertThat(dataset.listIndexes()).contains("body_idx"); + } + + assertQuery("SELECT count(*) FROM " + tableName, "SELECT 3"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + } + } + + @Test + public void testCreateIndexRequiresColumnArgument() + { + String tableName = "test_create_index_missing_col_" + System.currentTimeMillis(); + try { + assertUpdate("CREATE TABLE " + tableName + " (id bigint, body varchar)"); + + assertThatThrownBy(() -> getQueryRunner().execute( + "ALTER TABLE " + tableName + " EXECUTE create_index(index_type => 'fts')")) + .hasMessageContaining("column"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + } + } + + @Test + public void testCreateIndexRejectsUnsupportedIndexType() + { + String tableName = "test_create_index_bad_type_" + System.currentTimeMillis(); + try { + assertUpdate("CREATE TABLE " + tableName + " (id bigint, body varchar)"); + + assertThatThrownBy(() -> getQueryRunner().execute( + "ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'btree')")) + .hasMessageContaining("fts"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + } + } + + @Test + public void testOptimizeIndicesAfterInsertAndCompactWithDeferredRemap() + { + String tableName = "test_optimize_compact_" + System.currentTimeMillis(); + try { + assertUpdate("CREATE TABLE " + tableName + " (id bigint, body varchar)"); + assertUpdate("INSERT INTO " + tableName + " VALUES (1, 'hello world')", 1); + + getQueryRunner().execute("ALTER TABLE " + tableName + " EXECUTE create_index(column => 'body', index_type => 'fts')"); + + // Append more data after the index was built; the index now covers only part of the table. + assertUpdate("INSERT INTO " + tableName + " VALUES (2, 'goodbye world'), (3, 'another document')", 2); + assertQuery("SELECT count(*) FROM " + tableName, "SELECT 3"); + + // Incrementally catch up the index on the newly appended fragment, without a full rebuild. + getQueryRunner().execute("ALTER TABLE " + tableName + " EXECUTE optimize_indices()"); + assertQuery("SELECT count(*) FROM " + tableName, "SELECT 3"); + + // Compacting with defer_index_remap should not require rebuilding the index and should + // leave the table's data intact. + assertUpdate("INSERT INTO " + tableName + " VALUES (4, 'yet another document')", 1); + getQueryRunner().execute("ALTER TABLE " + tableName + " EXECUTE compact(defer_index_remap => true)"); + assertQuery("SELECT count(*) FROM " + tableName, "SELECT 4"); + + try (Dataset dataset = openDataset(tableName)) { + assertThat(dataset.listIndexes()).isNotEmpty(); + } + } + finally { + assertUpdate("DROP TABLE IF EXISTS " + tableName); + } + } +}